forestcalled/text-generation-webui
0
1import os2 3os.environ["WANDB_MODE"] = "offline"4# os.environ["WANDB_DISABLED"] = "true"5 6import json7import math8import random9import shutil10import sys11import threading12import time13import traceback14from datetime import datetime15from pathlib import Path16 17import gradio as gr18import torch19import transformers20from datasets import Dataset, load_dataset21from peft import (22 LoraConfig,23 get_peft_model,24 prepare_model_for_kbit_training,25 set_peft_model_state_dict26)27from peft.utils.other import \28 TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING as model_to_lora_modules29from transformers import is_torch_xpu_available30from transformers.models.auto.modeling_auto import (31 MODEL_FOR_CAUSAL_LM_MAPPING_NAMES32)33 34from modules import shared, ui, utils35from modules.evaluate import (36 calculate_perplexity,37 generate_markdown_table,38 save_past_evaluations39)40from modules.logging_colors import logger41from modules.models import reload_model42from modules.utils import natural_keys43 44MODEL_CLASSES = {v[1]: v[0] for v in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.items()}45PARAMETERS = ["lora_name", "always_override", "q_proj_en", "v_proj_en", "k_proj_en", "o_proj_en", "gate_proj_en", "down_proj_en", "up_proj_en", "save_steps", "micro_batch_size", "batch_size", "epochs", "learning_rate", "lr_scheduler_type", "lora_rank", "lora_alpha", "lora_dropout", "cutoff_len", "dataset", "eval_dataset", "format", "eval_steps", "raw_text_file", "overlap_len", "newline_favor_len", "higher_rank_limit", "warmup_steps", "optimizer", "hard_cut_string", "train_only_after", "stop_at_loss", "add_eos_token", "min_chars", "report_to"]46WANT_INTERRUPT = False47 48train_log = {}49train_template = {}50 51 52def create_ui():53 mu = shared.args.multi_user54 with gr.Tab("Training", elem_id="training-tab"):55 with gr.Tab('Train LoRA', elem_id='lora-train-tab'):56 tmp = gr.State('')57 with gr.Row():58 with gr.Column():59 gr.Markdown("[Tutorial](https://github.com/oobabooga/text-generation-webui/wiki/05-%E2%80%90-Training-Tab)")60 61 with gr.Row():62 copy_from = gr.Dropdown(label='Copy parameters from', value='None', choices=utils.get_available_loras(), elem_classes=['slim-dropdown'], interactive=not mu)63 ui.create_refresh_button(copy_from, lambda: None, lambda: {'choices': utils.get_available_loras()}, 'refresh-button', interactive=not mu)64 65 with gr.Row():66 with gr.Column(scale=5):67 lora_name = gr.Textbox(label='Name', info='The name of your new LoRA file')68 with gr.Column():69 always_override = gr.Checkbox(label='Override Existing Files', value=False, info='If the name is the same, checking will replace the existing file, and unchecking will load and continue from it (the rank must be the same).', elem_classes=['no-background'])70 71 with gr.Accordion(label='Target Modules', open=False):72 gr.Markdown("Selects which modules to target in training. Targeting more modules is closer to a full fine-tune at the cost of increased VRAM requirements and adapter size.\nNOTE: Only works for model_id='llama', other types will retain default training behavior and not use these settings.")73 with gr.Row():74 with gr.Column():75 q_proj_en = gr.Checkbox(label='Enable q_proj', value=True)76 with gr.Column():77 v_proj_en = gr.Checkbox(label='Enable v_proj', value=True)78 with gr.Column():79 k_proj_en = gr.Checkbox(label='Enable k_proj', value=False)80 with gr.Column():81 o_proj_en = gr.Checkbox(label='Enable o_proj', value=False)82 with gr.Column():83 gate_proj_en = gr.Checkbox(label='Enable gate_proj', value=False)84 with gr.Column():85 down_proj_en = gr.Checkbox(label='Enable down_proj', value=False)86 with gr.Column():87 up_proj_en = gr.Checkbox(label='Enable up_proj', value=False)88 89 with gr.Row():90 with gr.Column():91 lora_rank = gr.Slider(label='LoRA Rank', value=32, minimum=0, maximum=1024, step=4, info='Also called dimension count. Higher values = larger file, more content control. Smaller values = smaller file, less control. Use 4 or 8 for style, 128 or 256 to teach, 1024+ for fine-detail on big data. More VRAM is needed for higher ranks.')92 lora_alpha = gr.Slider(label='LoRA Alpha', value=64, minimum=0, maximum=2048, step=4, info='This divided by the rank becomes the scaling of the LoRA. Higher means stronger. A good standard value is twice your Rank.')93 batch_size = gr.Slider(label='Batch Size', value=128, minimum=0, maximum=1024, step=4, info='Global batch size. The two batch sizes together determine gradient accumulation (gradientAccum = batch / microBatch). Higher gradient accum values lead to better quality training.')94 micro_batch_size = gr.Slider(label='Micro Batch Size', value=4, minimum=1, maximum=128, step=1, info='Per-device batch size (NOTE: multiple devices not yet implemented). Increasing this will increase VRAM usage.')95 cutoff_len = gr.Slider(label='Cutoff Length', minimum=0, maximum=4096, value=256, step=32, info='Cutoff length for text input. Essentially, how long of a line of text to feed in at a time. Higher values require drastically more VRAM.')96 97 with gr.Column():98 save_steps = gr.Number(label='Save every n steps', value=0, info='If above 0, a checkpoint of the LoRA will be saved every time this many steps pass.')99 100 epochs = gr.Number(label='Epochs', value=3, info='Number of times every entry in the dataset should be fed into training. So 1 means feed each item in once, 5 means feed it in five times, etc.')101 learning_rate = gr.Textbox(label='Learning Rate', value='3e-4', info='In scientific notation. 3e-4 is a good starting base point. 1e-2 is extremely high, 1e-6 is extremely low.')102 with gr.Row():103 lr_scheduler_type = gr.Dropdown(label='LR Scheduler', value='linear', choices=['linear', 'constant', 'constant_with_warmup', 'cosine', 'cosine_with_restarts', 'polynomial', 'inverse_sqrt'], info='Learning rate scheduler - defines how the learning rate changes over time. "Constant" means never change, "linear" means to go in a straight line from the learning rate down to 0, cosine follows a curve, etc.', elem_classes=['slim-dropdown'])104 105 with gr.Accordion(label='Advanced Options', open=False):106 with gr.Row():107 with gr.Column():108 lora_dropout = gr.Slider(label='LoRA Dropout', minimum=0.0, maximum=1.0, step=0.025, value=0.05, info='Percentage probability for dropout of LoRA layers. This can help reduce overfitting. Most users should leave at default.')109 stop_at_loss = gr.Slider(label='Stop at loss', minimum=0.0, maximum=3.0, step=0.1, value=0.00, info='The process will automatically stop once the desired loss value is reached. (reasonable numbers are 1.5-1.8)')110 with gr.Row():111 optimizer = gr.Dropdown(label='Optimizer', value='adamw_torch', choices=['adamw_hf', 'adamw_torch', 'adamw_torch_fused', 'adamw_torch_xla', 'adamw_apex_fused', 'adafactor', 'adamw_bnb_8bit', 'adamw_anyprecision', 'sgd', 'adagrad'], info='Different optimizer implementation options, for advanced users. Effects of different options are not well documented yet.', elem_classes=['slim-dropdown'])112 113 with gr.Column():114 warmup_steps = gr.Number(label='Warmup Steps', value=100, info='For this many steps at the start, the learning rate will be lower than normal. This helps the trainer prepare the model and precompute statistics to improve the quality of training after the start.')115 train_only_after = gr.Textbox(label='Train Only After', value='', info='Only consider text *after* this string in any given chunk for training. For Alpaca datasets, use "### Response:" to only train the response and ignore the input.')116 117 add_eos_token = gr.Checkbox(label='Add EOS token', value=False, info="Adds EOS token for each dataset item. In case of raw text, the EOS will be added at the Hard Cut")118 119 higher_rank_limit = gr.Checkbox(label='Enable higher ranks', value=False, info='If checked, changes Rank/Alpha slider above to go much higher. This will not work without a datacenter-class GPU.')120 report_to = gr.Radio(label="Save detailed logs with", value="None", choices=["None", "wandb", "tensorboard"], interactive=True)121 122 with gr.Column():123 with gr.Tab(label='Formatted Dataset'):124 with gr.Row():125 format = gr.Dropdown(choices=utils.get_datasets('training/formats', 'json'), value='None', label='Data Format', info='The format file used to decide how to format the dataset input.', elem_classes=['slim-dropdown'], interactive=not mu)126 ui.create_refresh_button(format, lambda: None, lambda: {'choices': utils.get_datasets('training/formats', 'json')}, 'refresh-button', interactive=not mu)127 128 with gr.Row():129 dataset = gr.Dropdown(choices=utils.get_datasets('training/datasets', 'json'), value='None', label='Dataset', info='The dataset file to use for training.', elem_classes=['slim-dropdown'], interactive=not mu)130 ui.create_refresh_button(dataset, lambda: None, lambda: {'choices': utils.get_datasets('training/datasets', 'json')}, 'refresh-button', interactive=not mu)131 132 with gr.Row():133 eval_dataset = gr.Dropdown(choices=utils.get_datasets('training/datasets', 'json'), value='None', label='Evaluation Dataset', info='The (optional) dataset file used to evaluate the model after training.', elem_classes=['slim-dropdown'], interactive=not mu)134 ui.create_refresh_button(eval_dataset, lambda: None, lambda: {'choices': utils.get_datasets('training/datasets', 'json')}, 'refresh-button', interactive=not mu)135 136 eval_steps = gr.Number(label='Evaluate every n steps', value=100, info='If an evaluation dataset is given, test it every time this many steps pass.')137 138 with gr.Tab(label="Raw text file"):139 with gr.Row():140 raw_text_file = gr.Dropdown(choices=utils.get_datasets('training/datasets', 'txt'), value='None', label='Text file', info='The raw text file to use for training.', elem_classes=['slim-dropdown'], interactive=not mu)141 ui.create_refresh_button(raw_text_file, lambda: None, lambda: {'choices': utils.get_datasets('training/datasets', 'txt')}, 'refresh-button', interactive=not mu)142 143 with gr.Row():144 with gr.Column():145 overlap_len = gr.Slider(label='Overlap Length', minimum=0, maximum=512, value=128, step=16, info='How many tokens from the prior chunk of text to include into the next chunk. (The chunks themselves will be of a size determined by Cutoff Length). Setting overlap to exactly half the cutoff length may be ideal.')146 newline_favor_len = gr.Slider(label='Prefer Newline Cut Length', minimum=0, maximum=512, value=128, step=16, info='Length (in characters, not tokens) of the maximum distance to shift an overlap cut by to ensure chunks cut at newlines. If too low, cuts may occur in the middle of lines.')147 148 with gr.Column():149 hard_cut_string = gr.Textbox(label='Hard Cut String', value='\\n\\n\\n', info='String that indicates a hard cut between text parts. Helps prevent unwanted overlap.')150 min_chars = gr.Number(label='Ignore small blocks', value=0, info='Ignore Hard Cut blocks that have less or equal characters than this number')151 152 with gr.Row():153 start_button = gr.Button("Start LoRA Training", variant='primary', interactive=not mu)154 stop_button = gr.Button("Interrupt", interactive=not mu)155 156 output = gr.Markdown(value="Ready")157 158 with gr.Tab('Perplexity evaluation', elem_id='evaluate-tab'):159 with gr.Row():160 with gr.Column():161 models = gr.Dropdown(utils.get_available_models(), label='Models', multiselect=True, interactive=not mu)162 evaluate_text_file = gr.Dropdown(choices=['wikitext', 'ptb', 'ptb_new'] + utils.get_datasets('training/datasets', 'txt')[1:], value='wikitext', label='Input dataset', info='The raw text file on which the model will be evaluated. The first options are automatically downloaded: wikitext, ptb, and ptb_new. The next options are your local text files under training/datasets.', interactive=not mu)163 with gr.Row():164 with gr.Column():165 stride_length = gr.Slider(label='Stride', minimum=0, maximum=32768, value=512, step=256, info='Used to make the evaluation faster at the cost of accuracy. 1 = slowest but most accurate. 512 is a common value.')166 167 with gr.Column():168 max_length = gr.Slider(label='max_length', minimum=0, maximum=shared.settings['truncation_length_max'], value=0, step=256, info='The context for each evaluation. If set to 0, the maximum context length for the model will be used.')169 170 with gr.Row():171 start_current_evaluation = gr.Button("Evaluate loaded model", interactive=not mu)172 start_evaluation = gr.Button("Evaluate selected models", interactive=not mu)173 stop_evaluation = gr.Button("Interrupt", interactive=not mu)174 175 with gr.Column():176 evaluation_log = gr.Markdown(value='')177 178 evaluation_table = gr.Dataframe(value=generate_markdown_table(), interactive=True)179 with gr.Row():180 save_comments = gr.Button('Save comments', elem_classes="small-button", interactive=not mu)181 refresh_table = gr.Button('Refresh the table', elem_classes="small-button", interactive=not mu)182 183 # Training events184 all_params = [lora_name, always_override, q_proj_en, v_proj_en, k_proj_en, o_proj_en, gate_proj_en, down_proj_en, up_proj_en, save_steps, micro_batch_size, batch_size, epochs, learning_rate, lr_scheduler_type, lora_rank, lora_alpha, lora_dropout, cutoff_len, dataset, eval_dataset, format, eval_steps, raw_text_file, overlap_len, newline_favor_len, higher_rank_limit, warmup_steps, optimizer, hard_cut_string, train_only_after, stop_at_loss, add_eos_token, min_chars, report_to]185 186 copy_from.change(do_copy_params, [copy_from] + all_params, all_params)187 start_button.click(do_train, all_params, output)188 stop_button.click(do_interrupt, None, None, queue=False)189 higher_rank_limit.change(change_rank_limit, [higher_rank_limit], [lora_rank, lora_alpha])190 191 # Evaluation events. For some reason, the interrupt event192 # doesn't work with the .then() syntax, so I write them one193 # by one in this ugly but functional way.194 ev = start_evaluation.click(calculate_perplexity, [models, evaluate_text_file, stride_length, max_length], evaluation_log, show_progress=False)195 start_evaluation.click(generate_markdown_table, None, evaluation_table, show_progress=False)196 197 start_current_evaluation.click(lambda: ['current model'], None, tmp)198 ev_cur = start_current_evaluation.click(calculate_perplexity, [tmp, evaluate_text_file, stride_length, max_length], evaluation_log, show_progress=False)199 start_current_evaluation.click(generate_markdown_table, None, evaluation_table, show_progress=False)200 201 stop_evaluation.click(None, None, None, cancels=[ev, ev_cur], queue=False)202 refresh_table.click(generate_markdown_table, None, evaluation_table, show_progress=True)203 save_comments.click(204 save_past_evaluations, evaluation_table, None).then(205 lambda: "Comments saved.", None, evaluation_log, show_progress=False)206 207 208def do_interrupt():209 global WANT_INTERRUPT210 WANT_INTERRUPT = True211 212 213def do_copy_params(lora_name: str, *args):214 f_name = f"{shared.args.lora_dir}/{clean_path(None, lora_name)}/training_parameters.json"215 if Path(f_name).is_file():216 with open(f_name, 'r', encoding='utf-8') as format_file:217 params: dict[str, str] = json.load(format_file)218 else:219 params = {}220 221 result = list()222 for i in range(0, len(PARAMETERS)):223 key = PARAMETERS[i]224 if key in params:225 result.append(params[key])226 else:227 result.append(args[i])228 229 return result230 231 232def change_rank_limit(use_higher_ranks: bool):233 mult = 2 if use_higher_ranks else 1234 return {"maximum": 1024 * mult, "__type__": "update"}, {"maximum": 2048 * mult, "__type__": "update"}235 236 237def clean_path(base_path: str, path: str):238 """Strips unusual symbols and forcibly builds a path as relative to the intended directory."""239 path = path.replace('\\', '/').replace('..', '_')240 if base_path is None:241 return path242 243 return f'{Path(base_path).absolute()}/{path}'244 245 246def backup_adapter(input_folder):247 # Get the creation date of the file adapter_model.bin248 try:249 adapter_file = Path(f"{input_folder}/adapter_model.bin")250 if adapter_file.is_file():251 252 logger.info("Backing up existing LoRA adapter")253 creation_date = datetime.fromtimestamp(adapter_file.stat().st_ctime)254 creation_date_str = creation_date.strftime("Backup-%Y-%m-%d")255 256 # Create the new subfolder257 subfolder_path = Path(f"{input_folder}/{creation_date_str}")258 subfolder_path.mkdir(parents=True, exist_ok=True)259 260 # Check if the file already exists in the subfolder261 backup_adapter_file = Path(f"{input_folder}/{creation_date_str}/adapter_model.bin")262 if backup_adapter_file.is_file():263 print(" - Backup already exists. Skipping backup process.")264 return265 266 # Copy existing files to the new subfolder267 existing_files = Path(input_folder).iterdir()268 for file in existing_files:269 if file.is_file():270 shutil.copy2(file, subfolder_path)271 except Exception as e:272 print("An error occurred in backup_adapter:", str(e))273 274 275def calc_trainable_parameters(model):276 trainable_params = 0277 all_param = 0278 for _, param in model.named_parameters():279 num_params = param.numel()280 # if using DS Zero 3 and the weights are initialized empty281 if num_params == 0 and hasattr(param, "ds_numel"):282 num_params = param.ds_numel283 284 all_param += num_params285 if param.requires_grad:286 trainable_params += num_params287 288 return trainable_params, all_param289 290 291def do_train(lora_name: str, always_override: bool, q_proj_en: bool, v_proj_en: bool, k_proj_en: bool, o_proj_en: bool, gate_proj_en: bool, down_proj_en: bool, up_proj_en: bool, save_steps: int, micro_batch_size: int, batch_size: int, epochs: int, learning_rate: str, lr_scheduler_type: str, lora_rank: int, lora_alpha: int, lora_dropout: float, cutoff_len: int, dataset: str, eval_dataset: str, format: str, eval_steps: int, raw_text_file: str, overlap_len: int, newline_favor_len: int, higher_rank_limit: bool, warmup_steps: int, optimizer: str, hard_cut_string: str, train_only_after: str, stop_at_loss: float, add_eos_token: bool, min_chars: int, report_to: str):292 293 if shared.args.monkey_patch:294 from alpaca_lora_4bit.monkeypatch.peft_tuners_lora_monkey_patch import (295 replace_peft_model_with_int4_lora_model296 )297 replace_peft_model_with_int4_lora_model()298 299 global WANT_INTERRUPT300 WANT_INTERRUPT = False301 302 # == Input validation / processing ==303 yield "Preparing the input..."304 lora_file_path = clean_path(None, lora_name)305 if lora_file_path.strip() == '':306 yield "Missing or invalid LoRA file name input."307 return308 309 lora_file_path = f"{Path(shared.args.lora_dir)}/{lora_file_path}"310 actual_lr = float(learning_rate)311 model_type = type(shared.model).__name__312 313 if model_type in MODEL_CLASSES:314 model_id = MODEL_CLASSES[model_type]315 else:316 model_id = "llama"317 if model_type == "PeftModelForCausalLM":318 if len(shared.lora_names) > 0:319 yield "You are trying to train a LoRA while you already have another LoRA loaded. This will work, but may have unexpected effects. *(Will continue anyway in 5 seconds, press `Interrupt` to stop.)*"320 logger.warning("Training LoRA over top of another LoRA. May have unexpected effects.")321 else:322 yield "Model ID not matched due to LoRA loading. Consider reloading base model. *(Will continue anyway in 5 seconds, press `Interrupt` to stop.)*"323 logger.warning("Model ID not matched due to LoRA loading. Consider reloading base model.")324 else:325 yield "LoRA training has only currently been validated for LLaMA, OPT, GPT-J, and GPT-NeoX models. Unexpected errors may follow. *(Will continue anyway in 5 seconds, press `Interrupt` to stop.)*"326 logger.warning(f"LoRA training has only currently been validated for LLaMA, OPT, GPT-J, and GPT-NeoX models. (Found model type: {model_type})")327 328 time.sleep(5)329 330 if shared.args.loader == 'GPTQ-for-LLaMa' and not shared.args.monkey_patch:331 yield "LoRA training with GPTQ-for-LLaMa requires loading with `--monkey-patch`"332 return333 334 if cutoff_len <= 0 or micro_batch_size <= 0 or batch_size <= 0 or actual_lr <= 0 or lora_rank <= 0 or lora_alpha <= 0:335 yield "Cannot input zeroes."336 return337 338 gradient_accumulation_steps = batch_size // micro_batch_size339 shared.tokenizer.pad_token_id = 0340 shared.tokenizer.padding_side = "left"341 342 # Populate target_modules list with chosen X_proj modules. Llama-based models only atm, non-llama will revert to default behavior.343 def list_target_modules(model_id):344 if model_id != "llama":345 return model_to_lora_modules[model_id]346 347 available_modules = {348 "gate": gate_proj_en,349 "down": down_proj_en,350 "up": up_proj_en,351 "q": q_proj_en,352 "v": v_proj_en,353 "k": k_proj_en,354 "o": o_proj_en,355 }356 target_mods = [f"{name}_proj" for name, enabled in available_modules.items() if enabled]357 return target_mods358 359 def encode(text, add_bos_token):360 result = shared.tokenizer.encode(text, truncation=True, max_length=cutoff_len)361 # Check if the first two tokens are BOS362 if len(result) >= 2 and result[:2] == [shared.tokenizer.bos_token_id, shared.tokenizer.bos_token_id]:363 result = result[1:]364 365 if not add_bos_token and result[0] == shared.tokenizer.bos_token_id:366 result = result[1:]367 return result368 369 def tokenize(prompt, append_eos_token=False):370 371 if train_only_after == '' or train_only_after not in prompt:372 input_ids = encode(prompt, True)373 374 if append_eos_token and input_ids[-1] != shared.tokenizer.eos_token_id and len(input_ids) < cutoff_len:375 input_ids.append(shared.tokenizer.eos_token_id)376 377 input_ids = [shared.tokenizer.pad_token_id] * (cutoff_len - len(input_ids)) + input_ids378 labels = [1] * len(input_ids)379 380 else:381 ind = prompt.index(train_only_after) + len(train_only_after)382 before_tokens = encode(prompt[:ind], True)383 after_tokens = encode(prompt[ind:], False)384 385 if append_eos_token and after_tokens[-1] != shared.tokenizer.eos_token_id:386 after_tokens.append(shared.tokenizer.eos_token_id)387 388 full_length = len(after_tokens) + len(before_tokens)389 if full_length > cutoff_len:390 after_tokens = after_tokens[:cutoff_len - len(before_tokens)]391 else:392 before_tokens = [shared.tokenizer.pad_token_id] * (cutoff_len - full_length) + before_tokens393 394 input_ids = before_tokens + after_tokens395 labels = [-100] * len(before_tokens) + [1] * len(after_tokens)396 397 input_ids = torch.tensor(input_ids)398 return {399 "input_ids": input_ids,400 "labels": labels,401 "attention_mask": input_ids.ne(shared.tokenizer.pad_token_id),402 }403 404 train_template.clear()405 406 # == Prep the dataset, format, etc ==407 if raw_text_file not in ['None', '']:408 train_template["template_type"] = "raw_text"409 logger.info("Loading raw text file dataset")410 fullpath = clean_path('training/datasets', f'{raw_text_file}')411 fullpath = Path(fullpath)412 if fullpath.is_dir():413 logger.info('Training path directory {}'.format(raw_text_file))414 raw_text = ""415 file_paths = sorted(fullpath.glob('*.txt'), key=lambda path: natural_keys(path.name))416 for file_path in file_paths:417 if file_path.is_file():418 with file_path.open('r', encoding='utf-8') as file:419 raw_text += file.read().replace('\r', '')420 421 logger.info(f"Loaded training file: {file_path.name}")422 else:423 with open(clean_path('training/datasets', f'{raw_text_file}.txt'), 'r', encoding='utf-8') as file:424 raw_text = file.read().replace('\r', '')425 426 cut_string = hard_cut_string.replace('\\n', '\n')427 eos_added = 0428 out_tokens = []429 for text_part in raw_text.split(cut_string):430 if len(text_part.strip()) <= min_chars:431 continue432 433 tokens = shared.tokenizer.encode(text_part)434 if add_eos_token:435 tokens.append(shared.tokenizer.eos_token_id)436 eos_added += 1437 438 step = cutoff_len - overlap_len439 if step <= 0:440 yield f"Error: overlap_len ({overlap_len}) cannot be greater than or equal to cutoff_len ({cutoff_len})"441 return442 443 out_tokens.extend(split_chunks(tokens, cutoff_len, step))444 445 if eos_added > 0:446 print(f"EOS added to {eos_added} text blocks")447 448 del raw_text # Note: could be a gig for a large dataset, so delete redundant data as we go to be safe on RAM449 text_chunks = [shared.tokenizer.decode(x) for x in out_tokens]450 del out_tokens451 if newline_favor_len > 0:452 text_chunks = [cut_chunk_for_newline(x, newline_favor_len) for x in text_chunks]453 454 train_data = Dataset.from_list([tokenize(x) for x in text_chunks])455 del text_chunks456 eval_data = None457 else:458 if dataset in ['None', '']:459 yield "Missing dataset choice input, cannot continue."460 return461 462 if format in ['None', '']:463 yield "Missing format choice input, cannot continue."464 return465 466 train_template["template_type"] = "dataset"467 468 with open(clean_path('training/formats', f'{format}.json'), 'r', encoding='utf-8-sig') as formatFile:469 format_data: dict[str, str] = json.load(formatFile)470 471 # == store training prompt ==472 for _, value in format_data.items():473 prompt_key = f"template_{len(train_template)}"474 train_template[prompt_key] = value475 476 def generate_prompt(data_point: dict[str, str]):477 for options, data in format_data.items():478 if set(options.split(',')) == set(x[0] for x in data_point.items() if (type(x[1]) is str and len(x[1].strip()) > 0)):479 for key, val in data_point.items():480 if type(val) is str:481 data = data.replace(f'%{key}%', val)482 return data483 raise RuntimeError(f'Data-point "{data_point}" has no keyset match within format "{list(format_data.keys())}"')484 485 def generate_and_tokenize_prompt(data_point):486 prompt = generate_prompt(data_point)487 return tokenize(prompt, add_eos_token)488 489 logger.info("Loading JSON datasets")490 data = load_dataset("json", data_files=clean_path('training/datasets', f'{dataset}.json'))491 train_data = data['train'].map(generate_and_tokenize_prompt, new_fingerprint='%030x' % random.randrange(16**30))492 493 if eval_dataset == 'None':494 eval_data = None495 else:496 eval_data = load_dataset("json", data_files=clean_path('training/datasets', f'{eval_dataset}.json'))497 eval_data = eval_data['train'].map(generate_and_tokenize_prompt, new_fingerprint='%030x' % random.randrange(16**30))498 499 # == We MUST reload model if it went through any previous training, even failed one ==500 if shared.model_dirty_from_training:501 selected_model = shared.model_name502 if selected_model:503 print("\033[1;31;1m(Model has been modified by previous training, it needs to be reloaded...)\033[0;37;0m")504 try:505 yield f"Reloading {selected_model}..."506 reload_model()507 if shared.model is not None:508 print("Model reloaded OK, continue with training.")509 else:510 return f"Failed to load {selected_model}."511 except:512 exc = traceback.format_exc()513 logger.error('Failed to reload the model.')514 print(exc)515 return exc.replace('\n', '\n\n')516 517 # == Start prepping the model itself ==518 if not hasattr(shared.model, 'lm_head') or hasattr(shared.model.lm_head, 'weight'):519 logger.info("Getting model ready")520 prepare_model_for_kbit_training(shared.model)521 522 # base model is now frozen and should not be reused for any other LoRA training than this one523 shared.model_dirty_from_training = True524 525 logger.info("Preparing for training")526 config = LoraConfig(527 r=lora_rank,528 lora_alpha=lora_alpha,529 target_modules=list_target_modules(model_id),530 lora_dropout=lora_dropout,531 bias="none",532 task_type="CAUSAL_LM"533 )534 535 # == Backup the existing adapter ==536 if not always_override:537 backup_adapter(lora_file_path)538 539 # == get model trainable params540 model_trainable_params, model_all_params = calc_trainable_parameters(shared.model)541 542 try:543 logger.info("Creating LoRA model")544 lora_model = get_peft_model(shared.model, config)545 if not always_override and Path(f"{lora_file_path}/adapter_model.bin").is_file():546 logger.info("Loading existing LoRA data")547 state_dict_peft = torch.load(f"{lora_file_path}/adapter_model.bin", weights_only=True)548 set_peft_model_state_dict(lora_model, state_dict_peft)549 except:550 yield traceback.format_exc().replace('\n', '\n\n')551 return552 553 if shared.args.monkey_patch:554 from alpaca_lora_4bit.autograd_4bit import Autograd4bitQuantLinear555 from alpaca_lora_4bit.models import Linear4bitLt556 for _, m in lora_model.named_modules():557 if isinstance(m, Autograd4bitQuantLinear) or isinstance(m, Linear4bitLt):558 if m.is_v1_model:559 m.zeros = m.zeros.half()560 m.scales = m.scales.half()561 562 class Tracked():563 def __init__(self):564 self.current_steps = 0565 self.max_steps = 0566 self.did_save = False567 568 tracked = Tracked()569 actual_save_steps = math.ceil(save_steps / gradient_accumulation_steps)570 571 class Callbacks(transformers.TrainerCallback):572 def on_step_begin(self, args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs):573 tracked.current_steps = state.global_step * gradient_accumulation_steps574 tracked.max_steps = state.max_steps * gradient_accumulation_steps575 if WANT_INTERRUPT:576 control.should_epoch_stop = True577 control.should_training_stop = True578 elif state.global_step > 0 and actual_save_steps > 0 and state.global_step % actual_save_steps == 0:579 lora_model.save_pretrained(f"{lora_file_path}/checkpoint-{tracked.current_steps}/")580 # Save log581 with open(f"{lora_file_path}/checkpoint-{tracked.current_steps}/training_log.json", 'w', encoding='utf-8') as file:582 json.dump(train_log, file, indent=2)583 # == Save training prompt ==584 with open(f"{lora_file_path}/checkpoint-{tracked.current_steps}/training_prompt.json", 'w', encoding='utf-8') as file:585 json.dump(train_template, file, indent=2)586 587 def on_substep_end(self, args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs):588 tracked.current_steps += 1589 if WANT_INTERRUPT:590 control.should_epoch_stop = True591 control.should_training_stop = True592 593 def on_log(self, args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, logs, **kwargs):594 train_log.update(logs)595 train_log.update({"current_steps": tracked.current_steps})596 if WANT_INTERRUPT:597 print("\033[1;31;1mInterrupted by user\033[0;37;0m")598 599 print(f"\033[1;30;40mStep: {tracked.current_steps} \033[0;37;0m", end='')600 if 'loss' in logs:601 loss = float(logs['loss'])602 if loss <= stop_at_loss:603 control.should_epoch_stop = True604 control.should_training_stop = True605 print(f"\033[1;31;1mStop Loss {stop_at_loss} reached.\033[0;37;0m")606 607 trainer = transformers.Trainer(608 model=lora_model,609 train_dataset=train_data,610 eval_dataset=eval_data,611 args=transformers.TrainingArguments(612 report_to=report_to if report_to != "None" else None,613 per_device_train_batch_size=micro_batch_size,614 gradient_accumulation_steps=gradient_accumulation_steps,615 warmup_steps=math.ceil(warmup_steps / gradient_accumulation_steps),616 num_train_epochs=epochs,617 learning_rate=actual_lr,618 fp16=False if shared.args.cpu else True,619 optim=optimizer,620 logging_steps=2 if stop_at_loss > 0 else 5,621 evaluation_strategy="steps" if eval_data is not None else "no",622 eval_steps=math.ceil(eval_steps / gradient_accumulation_steps) if eval_data is not None else None,623 save_strategy="steps" if eval_data is not None else "no",624 output_dir=lora_file_path,625 lr_scheduler_type=lr_scheduler_type,626 load_best_model_at_end=eval_data is not None,627 # TODO: Enable multi-device support628 ddp_find_unused_parameters=None,629 no_cuda=shared.args.cpu,630 use_ipex=True if is_torch_xpu_available and not shared.args.cpu else False631 ),632 data_collator=transformers.DataCollatorForLanguageModeling(shared.tokenizer, mlm=False),633 callbacks=list([Callbacks()])634 )635 636 lora_model.config.use_cache = False637 638 if torch.__version__ >= "2" and sys.platform != "win32":639 lora_model = torch.compile(lora_model)640 641 # == Save parameters for reuse ==642 with open(f"{lora_file_path}/training_parameters.json", 'w', encoding='utf-8') as file:643 vars = locals()644 json.dump({x: vars[x] for x in PARAMETERS}, file, indent=2)645 646 # == Save training prompt ==647 with open(f"{lora_file_path}/training_prompt.json", 'w', encoding='utf-8') as file:648 json.dump(train_template, file, indent=2)649 650 # == Main run and monitor loop ==651 logger.info("Starting training")652 yield "Starting..."653 654 lora_trainable_param, lora_all_param = calc_trainable_parameters(lora_model)655 656 projections_string = ", ".join([projection.replace("_proj", "") for projection in list_target_modules(model_id)])657 658 print(f"Training '{model_id}' model using ({projections_string}) projections")659 660 if lora_all_param > 0:661 print(f"Trainable params: {lora_trainable_param:,d} ({100 * lora_trainable_param / lora_all_param:.4f} %), All params: {lora_all_param:,d} (Model: {model_all_params:,d})")662 663 train_log.update({"base_model_name": shared.model_name})664 train_log.update({"base_model_class": shared.model.__class__.__name__})665 train_log.update({"base_loaded_in_4bit": getattr(lora_model, "is_loaded_in_4bit", False)})666 train_log.update({"base_loaded_in_8bit": getattr(lora_model, "is_loaded_in_8bit", False)})667 train_log.update({"projections": projections_string})668 669 if stop_at_loss > 0:670 print(f"Monitoring loss \033[1;31;1m(Auto-Stop at: {stop_at_loss})\033[0;37;0m")671 672 if WANT_INTERRUPT:673 yield "Interrupted before start."674 return675 676 def log_train_dataset(trainer):677 decoded_entries = []678 # Try to decode the entries and write the log file679 try:680 # Iterate over the first 10 elements in the dataset (or fewer if there are less than 10)681 for i in range(min(10, len(trainer.train_dataset))):682 decoded_text = shared.tokenizer.decode(trainer.train_dataset[i]['input_ids'])683 decoded_entries.append({"value": decoded_text})684 685 # Write the log file686 Path('logs').mkdir(exist_ok=True)687 with open(Path('logs/train_dataset_sample.json'), 'w') as json_file:688 json.dump(decoded_entries, json_file, indent=4)689 690 logger.info("Log file 'train_dataset_sample.json' created in the 'logs' directory.")691 except Exception as e:692 logger.error(f"Failed to create log file due to error: {e}")693 694 def threaded_run():695 log_train_dataset(trainer)696 trainer.train()697 # Note: save in the thread in case the gradio thread breaks (eg browser closed)698 lora_model.save_pretrained(lora_file_path)699 logger.info("LoRA training run is completed and saved.")700 # Save log701 with open(f"{lora_file_path}/training_log.json", 'w', encoding='utf-8') as file:702 json.dump(train_log, file, indent=2)703 704 thread = threading.Thread(target=threaded_run)705 thread.start()706 last_step = 0707 start_time = time.perf_counter()708 709 while thread.is_alive():710 time.sleep(0.5)711 if WANT_INTERRUPT:712 yield "Interrupting, please wait... *(Run will stop after the current training step completes.)*"713 714 elif tracked.current_steps != last_step:715 last_step = tracked.current_steps716 time_elapsed = time.perf_counter() - start_time717 if time_elapsed <= 0:718 timer_info = ""719 total_time_estimate = 999720 else:721 its = tracked.current_steps / time_elapsed722 if its > 1:723 timer_info = f"`{its:.2f}` it/s"724 else:725 timer_info = f"`{1.0/its:.2f}` s/it"726 727 total_time_estimate = (1.0 / its) * (tracked.max_steps)728 729 yield f"Running... **{tracked.current_steps}** / **{tracked.max_steps}** ... {timer_info}, {format_time(time_elapsed)} / {format_time(total_time_estimate)} ... {format_time(total_time_estimate - time_elapsed)} remaining"730 731 # Saving in the train thread might fail if an error occurs, so save here if so.732 if not tracked.did_save:733 logger.info("Training complete, saving")734 lora_model.save_pretrained(lora_file_path)735 736 if WANT_INTERRUPT:737 logger.info("Training interrupted.")738 yield f"Interrupted. Incomplete LoRA saved to `{lora_file_path}`."739 else:740 logger.info("Training complete!")741 yield f"Done! LoRA saved to `{lora_file_path}`.\n\nBefore testing your new LoRA, make sure to first reload the model, as it is currently dirty from training."742 743 744def split_chunks(arr, size, step):745 for i in range(0, len(arr), step):746 yield arr[i:i + size]747 748 749def cut_chunk_for_newline(chunk: str, max_length: int):750 if '\n' not in chunk:751 return chunk752 753 first_newline = chunk.index('\n')754 if first_newline < max_length:755 chunk = chunk[first_newline + 1:]756 757 if '\n' not in chunk:758 return chunk759 760 last_newline = chunk.rindex('\n')761 if len(chunk) - last_newline < max_length:762 chunk = chunk[:last_newline]763 764 return chunk765 766 767def format_time(seconds: float):768 if seconds < 120:769 return f"`{seconds:.0f}` seconds"770 771 minutes = seconds / 60772 if minutes < 120:773 return f"`{minutes:.0f}` minutes"774 775 hours = minutes / 60776 return f"`{hours:.0f}` hours"777