diff --git a/.gitignore b/.gitignore index 3630a70..737d8ef 100644 --- a/.gitignore +++ b/.gitignore @@ -6,7 +6,10 @@ models/ results/ datasets/ llm-comparator/ -/data/ +# Ignore data files but not scripts +/data/processed_datasets/ +/data/raw_datasets/ +/data/distil/ __pycache__/ *egg-info/ .vscode/ @@ -22,7 +25,6 @@ plots/ *.safetensors *tfevents* watcher_state.yaml -ds_config.json eval_results/ .github/ *.code-workspace diff --git a/README.md b/README.md index 2b11eb5..2f0dc8e 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,11 @@ print(model.decode(outputs)) ```bash WANDB_MODE=disabled run uv run intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=1 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --load_best_model_at_end=False --add_negative_prompt=False --add_repeat_prompt=False --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 ``` + +KL loss +```bash + WANDB_MODE=disabled run uv run intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=1 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --load_best_model_at_end=False --add_negative_prompt=False --add_repeat_prompt=False --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.0 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=True --eval_steps=99999999 --use_liger_kernel=False + ``` ### FWQA-v2 Level-0 Tiny ```bash WANDB_MODE=disabled run uv run intx_sft.py configs/fw_qa_v2_level_0_tiny.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=64 --gradient_accumulation_steps=8 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=1 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --max_packed_inp_len=24000 --max_packed_ctx_len=48000 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=50 @@ -71,6 +76,12 @@ uv run python data/generate_fav_num.py uv run python data/generate_fav_num_big.py ``` +### SQuAD +```bash +# for some reason download directly through `load_dataset` does not work +HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download --repo-type dataset rajpurkar/squad --local-dir data/raw_datasets/squad +``` + ### Data Generation (v2) ***Download a subset of generated Fineweb-QA directly*** ```bash diff --git a/chat_templates/google/gemma-2-2b-it.jinja b/chat_templates/google/gemma-2-2b-it.jinja index 23129b0..4b683e5 100644 --- a/chat_templates/google/gemma-2-2b-it.jinja +++ b/chat_templates/google/gemma-2-2b-it.jinja @@ -14,7 +14,7 @@ {%- if message['role'] == 'user' and loop.first and system_message is defined %} {{ '' + role + '\n' + system_message + '\n\n' + message['content'] | trim + '\n' }} {%- elif message['role'] == 'assistant' %} - {{ '' + role + '\n' }}{% generation %}{{ message['content'] | trim + '\n' }}{% endgeneration %} + {{ '' + role + '\n' }}{% generation %}{{ message['content'] + '' }}{% endgeneration %}{{ '\n' }} {%- else %} {{ '' + role + '\n' + message['content'] | trim + '\n' }} {%- endif %} diff --git a/chat_templates/google/gemma-2-2b-it.jinja.bak(w_trim) b/chat_templates/google/gemma-2-2b-it.jinja.bak(w_trim) new file mode 100644 index 0000000..23129b0 --- /dev/null +++ b/chat_templates/google/gemma-2-2b-it.jinja.bak(w_trim) @@ -0,0 +1,24 @@ +{{- bos_token }} +{%- if messages[0]['role'] == 'system' %} + {%- set system_message = messages[0]['content'] %} + {%- set loop_messages = messages[1:] %} +{%- else %} + {%- set loop_messages = messages %} +{%- endif %} +{% for message in loop_messages %} + {% if (message['role'] == 'assistant') %} + {% set role = 'model' %} + {% else %} + {% set role = message['role'] %} + {% endif %} + {%- if message['role'] == 'user' and loop.first and system_message is defined %} + {{ '' + role + '\n' + system_message + '\n\n' + message['content'] | trim + '\n' }} + {%- elif message['role'] == 'assistant' %} + {{ '' + role + '\n' }}{% generation %}{{ message['content'] | trim + '\n' }}{% endgeneration %} + {%- else %} + {{ '' + role + '\n' + message['content'] | trim + '\n' }} + {%- endif %} +{% endfor %} +{% if add_generation_prompt %} + {{'model\n'}} +{% endif %} diff --git a/configs/context_numbers_10_self_gen.yaml b/configs/context_numbers_10_self_gen.yaml new file mode 100644 index 0000000..566ad4b --- /dev/null +++ b/configs/context_numbers_10_self_gen.yaml @@ -0,0 +1,43 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: meta-llama/Llama-3.2-1B-Instruct +label_names: ["labels"] +add_repeat_prompt: false +add_negative_prompt: false +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 64 +per_device_eval_batch_size: 128 +max_new_tokens: 64 +gen_per_device_eval_batch_size: 128 +max_val_samples_per_ds: 500 +# optim: schedule_free_adamw +learning_rate: 0.0001 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +# LoRA +lora_r: 16 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: +- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/context_numbers_2_10 + +val_ds_names: +- context_numbers_2_10 diff --git a/configs/pwc_tiny.yaml b/configs/pwc_tiny.yaml new file mode 100644 index 0000000..136adcc --- /dev/null +++ b/configs/pwc_tiny.yaml @@ -0,0 +1,45 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: google/gemma-3-1b-it +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00004 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +warmup_steps: 100 + +dataloader_prefetch_factor: 16 +dataloader_num_workers: 8 + +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: + - pwc_tiny + +val_ds_names: + - pwc diff --git a/configs/qa_short_ctx_self_gen_no_fw_qa.yaml b/configs/qa_short_ctx_self_gen_no_fw_qa.yaml index 030e683..28bef65 100644 --- a/configs/qa_short_ctx_self_gen_no_fw_qa.yaml +++ b/configs/qa_short_ctx_self_gen_no_fw_qa.yaml @@ -39,10 +39,10 @@ target_modules: # data train_ds_names: - - self_gen/google/gemma-2-2b-it_temp_0.0/squad_compact - - self_gen/google/gemma-2-2b-it_temp_0.0/pwc_compact - - self_gen/google/gemma-2-2b-it_temp_0.0/ropes_compact - - self_gen/google/gemma-2-2b-it_temp_0.0/drop_compact + - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact + - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact + - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact + - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact val_ds_names: - squad diff --git a/configs/qa_short_ctx_self_gen_squad.yaml b/configs/qa_short_ctx_self_gen_squad.yaml new file mode 100644 index 0000000..27f2d0e --- /dev/null +++ b/configs/qa_short_ctx_self_gen_squad.yaml @@ -0,0 +1,45 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: google/gemma-3-1b-it +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00004 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +warmup_steps: 100 + +dataloader_prefetch_factor: 16 +dataloader_num_workers: 8 + +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: + - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact + +val_ds_names: + - squad diff --git a/data/build_squad_compact.py b/data/build_squad_compact.py index 9594f52..62b0f72 100644 --- a/data/build_squad_compact.py +++ b/data/build_squad_compact.py @@ -6,7 +6,7 @@ from tqdm import tqdm from ctx_to_lora.data.processing import closed_qa_prompting if __name__ == "__main__": - ds_name = "rajpurkar/squad" + ds_name = "data/raw_datasets/squad" for split in ["train", "validation"]: ctx_qa_dict = dict() diff --git a/data/download_self_gen_qa_logprobs.py b/data/download_self_gen_qa_logprobs.py new file mode 100644 index 0000000..1abcb25 --- /dev/null +++ b/data/download_self_gen_qa_logprobs.py @@ -0,0 +1,9 @@ +from huggingface_hub import snapshot_download + +if __name__ == "__main__": + fw_dir = "./data/raw_datasets/self_gen/" + snapshot_download( + "Rujikorn/self_gen_qa_logprobs", + repo_type="dataset", + local_dir=fw_dir, + ) diff --git a/data/prepare_distillation_dataset.py b/data/prepare_distillation_dataset.py new file mode 100644 index 0000000..d2a11e0 --- /dev/null +++ b/data/prepare_distillation_dataset.py @@ -0,0 +1,143 @@ +import argparse + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from ctx_to_lora.data.processing import load_and_process_dataset + + +def add_logits(row, model, tok, device, args): + ctx = row["context"] + qs = row["prompts"] + ans = row["responses"] + responses_masks = [] + all_logits_values = [] + all_logits_positions = [] + + for question, answer in zip(qs, ans): + messages = [ + { + "role": "user", + "content": f"Answer the question based on the following information:\n{ctx}.\n\n{question}", + }, + {"role": "assistant", "content": answer}, + ] + all_tokens = tok.apply_chat_template( + messages, + tokenize=True, + add_special_tokens=False, + truncation=False, + padding=False, + add_generation_prompt=False, + return_dict=True, + ) + + input = torch.tensor([all_tokens["input_ids"]], device=device) + with torch.no_grad(): + logits = model(input).logits + + top_logits, top_logits_positions = torch.topk( + logits, args.max_logits_to_store, dim=-1 + ) + + # Convert logits to specified precision + dtype_map = { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + "float16": torch.float16, + } + target_dtype = dtype_map[args.logit_precision] + top_logits = top_logits.to(target_dtype) + + tokens_to_mask = tok.apply_chat_template( + [messages[0]], # Just the user message + tokenize=True, + add_special_tokens=False, + truncation=False, + padding=False, + add_generation_prompt=True, + return_dict=True, + ) + number_of_tokens_to_mask = len(tokens_to_mask["input_ids"]) + + # Create mask: False for prompt tokens, True for response tokens + responses_mask = torch.zeros(len(all_tokens["input_ids"]), dtype=torch.bool) + responses_mask[number_of_tokens_to_mask:] = True + + responses_masks.append(responses_mask) + all_logits_values.append(top_logits.cpu()) + all_logits_positions.append(top_logits_positions.cpu()) + + row["labels_logits_values"] = all_logits_values + row["labels_logits_positions"] = all_logits_positions + row["responses_mask"] = responses_masks + return row + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--ds_name", required=True, help="name of the dataset") + ap.add_argument("--split", required=True, help="split of the dataset") + ap.add_argument( + "--max_logits_to_store", + type=int, + default=100, + help="max number of logits kept per completion", + ) + ap.add_argument( + "--limit_samples", + type=int, + default=None, + help="process only the first n dataset rows", + ) + ap.add_argument( + "--logit_precision", + default="bfloat16", + choices=["float32", "bfloat16", "float16"], + help="dtype for saved logits", + ) + ap.add_argument( + "--model_name", required=True, help="model name to use for generating logits" + ) + ap.add_argument( + "--output_path", required=True, help="path to save the processed dataset" + ) + args = ap.parse_args() + + # Load model and tokenizer + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = AutoModelForCausalLM.from_pretrained(args.model_name) + model = model.to(device) + model.eval() + tok = AutoTokenizer.from_pretrained(args.model_name) + + # Load dataset + ds = load_and_process_dataset( + args.ds_name, + args.split, + add_negative_prompt=False, + add_repeat_prompt=False, + repeat_prob=0, + is_pretrain=False, + streaming=False, + num_proc=8, + ) + + # Limit samples if specified + if args.limit_samples: + ds = ds.select(range(min(args.limit_samples, len(ds)))) + + # Process dataset with logits + ds = ds.map( + lambda row: add_logits(row, model, tok, device, args), + batched=False, + desc="Adding logits to dataset", + ) + + # Save the processed dataset + ds.save_to_disk(args.output_path) + print(f"Dataset saved to {args.output_path}") + + +if __name__ == "__main__": + main() diff --git a/data/self_generate_qa.py b/data/self_generate_qa.py index 9083e79..bae8205 100644 --- a/data/self_generate_qa.py +++ b/data/self_generate_qa.py @@ -4,7 +4,7 @@ import random import re from glob import glob -import pandas as pd +import numpy as np import yaml from datasets import Dataset, load_dataset from vllm import LLM, SamplingParams @@ -13,12 +13,15 @@ from ctx_to_lora.data.definitions import ( CLOSED_QA_INTX_TEMPLATES, RAW_DATA_DIR, SELF_GEN_DATA_DIR, + SELF_GEN_SYSTEM_MSG, ) from ctx_to_lora.data.processing import ( filter_none, get_preprocessing_fn, load_and_process_dataset, + tokenize_ctx_text, ) +from ctx_to_lora.model_loading import get_tokenizer from ctx_to_lora.utils import clear_gpu STOP_STRINGS = { @@ -31,14 +34,55 @@ STOP_STRINGS = { SYSTEM_TEMPLATE = ( "### SYSTEM INSTRUCTION ###\n" "You are an honest and helpful assistant.\n" - "You must use the information provided in the context for responding to the question.\n" + # "You must use the information provided in the context for responding to the question.\n" # "If the context does not contain enough information to answer the question, you must say so.\n" # "If that is the case, you can answer based on your knowledge, but you must clearly state that you are doing so.\n" "**DO NOT** hallucinate or make up information.\n" "### END OF SYSTEM INSTRUCTION ###" ) -PROMPT_TEMPLATE = "### Context ###\n{context}\n\n\n### Question ###\n{question}" +SELF_QA_INTX = ( + "---\n\n" + "# System Instruction\n" + "- The information provided is up-to-date information.\n" + "- When the provided information is not relevant to the question, ***ignore*** it and answer the question based on your knowledge.\n" + "- If the provided information is related to the question, incorporate it in your response.\n" + "\n---\n\n" + "# User\n" +) + +PRE_CTX = "# Information\n" + +PROMPT_TEMPLATE = ( + PRE_CTX + + "{context}\n\n---\n\n" + + + # "# System Instruction\n" + # "- The information provided is up-to-date information.\n" + # "- When the provided information is not relevant to the question, ***ignore*** it and answer the question based on your knowledge.\n" + # "- If the provided information is related to the question, incorporate it in your response.\n" + # "\n---\n\n" + # "# User\n{question}" + SELF_QA_INTX + + "{question}" +) + + +""" + +--- + +# System Instruction +- The information provided is up-to-date information. +- When the provided information is not relevant to the question, ***ignore*** it and answer the question based on your knowledge. +- If the provided information is related to the question, incorporate it in your response. + +--- +# User + +What team is Mbappe playing for? + +""" MODEL_CTX_LEN = { "google/gemma-2-27b-it": 8192, @@ -64,12 +108,12 @@ def load_config(config_path: str) -> dict: return config -def check_should_skip(txt: str, vllm_model: str) -> bool: - """Check if the response should be skipped based on stop strings.""" - for stop in STOP_STRINGS[vllm_model]: - if stop in txt[-len(stop) :]: - return (txt.split(stop)[0], False) # Found a valid stop string - return (txt, True) # No valid stop string found, skip this response +# def check_should_skip(txt: str, vllm_model: str) -> bool: +# """Check if the response should be skipped based on stop strings.""" +# for stop in STOP_STRINGS[vllm_model]: +# if stop in txt[-len(stop) :]: +# return (txt.split(stop)[0], False) # Found a valid stop string +# return (txt, True) # No valid stop string found, skip this response def get_dataset_configs( @@ -123,27 +167,27 @@ def create_messages( ctxs: list[str], questions: list[list[str]], vllm_model: str, system_template: str ) -> list[list[dict]]: """Create chat messages for the model.""" - if "gemma" in vllm_model: - # gemma models do not support system messages - return [ - [ - { - "role": "user", - "content": system_template + "\n\n\n" + get_prompt(ctx, q), - } - ] - for ctx, q_list in zip(ctxs, questions) - for q in q_list - ] - else: - return [ - [ - {"role": "system", "content": system_template}, - {"role": "user", "content": get_prompt(ctx, q)}, - ] - for ctx, q_list in zip(ctxs, questions) - for q in q_list + # if "gemma" in vllm_model: + # gemma models do not support system messages + return [ + [ + { + "role": "user", + "content": system_template + "\n\n\n" + get_prompt(ctx, q), + } ] + for ctx, q_list in zip(ctxs, questions) + for q in q_list + ] + # else: + # return [ + # [ + # {"role": "system", "content": system_template}, + # {"role": "user", "content": get_prompt(ctx, q)}, + # ] + # for ctx, q_list in zip(ctxs, questions) + # for q in q_list + # ] def self_generate( @@ -219,6 +263,25 @@ def self_generate( ds = ds.filter(filter_none, batched=False, num_proc=8) + tk = get_tokenizer(args.vllm_model, train=True) + + self_qa_intx_tokens = tk(SELF_QA_INTX, add_special_tokens=False)["input_ids"] + n_self_qa_intx_tokens = len(self_qa_intx_tokens) + pre_ctx_tokens = tk(PRE_CTX, add_special_tokens=False)["input_ids"] + n_pre_ctx_tokens = len(pre_ctx_tokens) + sys_tokens = tk(system_template.split("\n")[0], add_special_tokens=False)[ + "input_ids" + ] + n_sys_tokens = len(sys_tokens) + os.environ["TOKENIZERS_PARALLELISM"] = "true" + ds = ds.map( + tokenize_ctx_text, + fn_kwargs={"tokenizer": tk}, + batched=True, + batch_size=50_000, + keep_in_memory=True, + ) + ctxs = [sample["context"] for sample in ds] questions = [ [add_closed_qa_prompt(q, closed_qa_prob) for q in sample["prompts"] if q] @@ -226,59 +289,212 @@ def self_generate( ] print(f"Loaded {len(ctxs)} contexts and {len(questions)} questions") - messages = create_messages(ctxs, questions, args.vllm_model, SYSTEM_TEMPLATE) + k = 16 + fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}/{ds_name}/{split}/ds{shard_name}" - print(f"Generating from {len(messages)} contexts") - # TODO (distillation): make vllm outputs logits here too + chunk_size = 1_000 + for chunk_idx, start in enumerate(range(0, len(ctxs), chunk_size)): + print(f"Processing chunk {chunk_idx}") + + chunk_ctxs = ctxs[start : start + chunk_size] + chunk_questions = questions[start : start + chunk_size] + chunk_messages = create_messages( + chunk_ctxs, chunk_questions, args.vllm_model, SELF_GEN_SYSTEM_MSG + ) + print(f"Generating from {len(chunk_messages)} contexts") + + # Clear GPU memory before processing the next chunk + clear_gpu() + execute_qa_generation( + fpath + f"_{chunk_idx:04d}", + args, + llm, + temp, + tk, + self_qa_intx_tokens, + n_self_qa_intx_tokens, + sys_tokens, + n_sys_tokens, + chunk_ctxs, + ds["ctx_ids"][start : start + chunk_size], + chunk_questions, + chunk_messages, + k, + ) + + +def execute_qa_generation( + fpath, + args, + llm, + temp, + tk, + self_qa_intx_tokens, + n_self_qa_intx_tokens, + sys_tokens, + n_sys_tokens, + ctxs, + ctx_ids, + questions, + messages, + k, +): completions = llm.chat( messages, sampling_params=SamplingParams( max_tokens=1024, + logprobs=k, temperature=temp, - # needed for checking if stop tokens are present + seed=42, + spaces_between_special_tokens=False, skip_special_tokens=False, include_stop_str_in_output=True, ), ) - self_gen_data = {ctx: {"prompts": [], "responses": []} for ctx in ctxs} + self_gen_data = { + ctx: { + "ctx_ids": ctx_ids, + "input_ids": [], + "response_start_end": [], + "logprobs_vals": [], + "logprobs_indices": [], + } + for ctx, ctx_ids in zip(ctxs, ctx_ids) + } c = 0 n_skips = 0 + sys_start = None for ctx, q_list in zip(ctxs, questions): - for i, q in enumerate(q_list): - response = completions[c + i].outputs[0].text - response, skip = check_should_skip(response, args.vllm_model) - # skip = True - # for stop in STOP_STRINGS[args.vllm_model]: - # if stop == response[-len(stop) :]: - # # Check if response ends with stop string - # response = response.split(stop)[0] - # skip = False - # break - if skip: - print(f"Skipping due to missing stop string") + # self_gen_data[ctx]["ctx_ids"] = ctx_ids + for i, _ in enumerate(q_list): + # response = completions[c + i].outputs[0].text + reason = completions[c + i].outputs[0].finish_reason + if reason != "stop": + # print(f"idx: {c + i}") + print(f"finish_reason: {completions[c + i].outputs[0].finish_reason}") + print(f"Skipping due to finish_reason={reason} != 'stop'") n_skips += 1 continue - self_gen_data[ctx]["prompts"].append(q) - self_gen_data[ctx]["responses"].append(response) + # includes the logprob before the first response token + # but excludes the logprob from eos token + logp = completions[c + i].outputs[0].logprobs + + # len = num response tokens + n_response_tokens = len(completions[c + i].outputs[0].token_ids) + + logp_indices = np.empty((n_response_tokens, k), dtype=np.int32) + # float-16 is better for this range + logp_vals = np.empty((n_response_tokens, k), dtype=np.float16) + assert len(logp) == n_response_tokens, ( + f"Expected {n_response_tokens} logp entries, got {len(logp)}" + ) + + for li, info_d in enumerate(logp): + for j, (idx, tok_info) in enumerate(info_d.items()): + logp_indices[li, j] = idx + logp_vals[li, j] = tok_info.logprob + + prompt_ids = completions[c + i].prompt_token_ids # 1d list + # token_ids only includes generated tokens, not the prompt + response_token_ids = completions[c + i].outputs[0].token_ids # 1d list + all_ids = prompt_ids + response_token_ids + res_start = len(prompt_ids) + res_end = res_start + n_response_tokens + + if sys_start is None: + for ii in range(len(prompt_ids) - n_sys_tokens): + if prompt_ids[ii : ii + n_sys_tokens] == sys_tokens: + # found the start of the system message + sys_start = ii + break + + q_start = None + for ii in range( + len(prompt_ids) - n_self_qa_intx_tokens, + -1, + -1, + ): + if prompt_ids[ii : ii + n_self_qa_intx_tokens] == self_qa_intx_tokens: + # found the start of the user input + q_start = ii + n_self_qa_intx_tokens + break + + # bos + question + eos + start model turn + response + eos + input_ids = all_ids[:sys_start] + all_ids[q_start:res_end] + + # TODO: save prompt logprobs with qa_start + # TODO: also save the prompt tokens for output alignment during training + # prompt_logp = completions[c + i].prompt_logprobs + # q_len = res_start - q_start + # prompt_logp_indices = np.empty((q_len, k), dtype=np.int32) + # # float-16 is better for this range + # prompt_logp_vals = np.empty((q_len, k), dtype=np.float16) + # for li, info_d in enumerate(prompt_logp[q_start:res_start]): + # for j, (idx, tok_info) in enumerate(info_d.items()): + # # vllm returns logprob for gt token first, + # # which means that we might have k + 1 logprobs + # if j >= k: + # continue + # prompt_logp_indices[li, j] = idx + # prompt_logp_vals[li, j] = tok_info.logprob + + # relative to the input_ids + res_start = res_start - q_start + sys_start + res_end = res_start + n_response_tokens + + # arrays will be saved as nested lists of numbers + + self_gen_data[ctx]["input_ids"].append(input_ids) + # assume single-turn chat + self_gen_data[ctx]["response_start_end"].append((res_start, res_end)) + self_gen_data[ctx]["logprobs_vals"].append(logp_vals) + self_gen_data[ctx]["logprobs_indices"].append(logp_indices) + # NOTE: also have to shift back 1 for training + # self_gen_data[ctx]["prompt_start_end"].append((sys_start, res_start)) + # self_gen_data[ctx]["prompt_logprobs_vals"].append(prompt_logp_vals) + # self_gen_data[ctx]["prompt_logprobs_indices"].append(prompt_logp_indices) + c += i + 1 print(f"Skipped {n_skips} responses due to missing stop strings") samples = [ { - "context": ctx, - "prompts": q_list, - "responses": self_gen_data[ctx]["responses"], + # "context": ctx, + # "prompts": q_list, + # "responses": self_gen_data[ctx]["responses"], + "ctx_ids": self_gen_data[ctx]["ctx_ids"], + "input_ids": self_gen_data[ctx]["input_ids"], + "response_start_end": self_gen_data[ctx]["response_start_end"], + # "prompt_start_end": self_gen_data[ctx]["prompt_start_end"], + "logprobs_vals": self_gen_data[ctx]["logprobs_vals"], + "logprobs_indices": self_gen_data[ctx]["logprobs_indices"], } for ctx, q_list in zip(ctxs, questions) ] if args.debug: for sample in samples: - print(f"context={sample['context']}") - print(f"prompt={sample['prompts']}") - print(f"response={sample['responses']}") + print(f"context={tk.decode(sample['ctx_ids'])}") + print(f"QA={[tk.decode(ids) for ids in sample['input_ids']]}") + + for input_ids, (start, end) in zip( + sample["input_ids"], sample["response_start_end"] + ): + print(f"start={start}, end={end}") + print(f"response={tk.decode(input_ids[start:end])}") + + for input_ids, (start, end) in zip( + sample["input_ids"], sample["prompt_start_end"] + ): + print(f"start={start}, end={end}") + print(f"prompt={tk.decode(input_ids[start:end])}") + + print(f"logprobs_vals={[x.shape for x in sample['logprobs_vals']]}") + print(f"logprobs_indices={[x.shape for x in sample['logprobs_indices']]}") + for indices in sample["logprobs_indices"]: + print(f"logprobs_indices={indices[-1]}") print("=" * 80) print(f"Generated {len(samples)} samples") @@ -286,16 +502,18 @@ def self_generate( # Save results if not args.debug: - df = pd.DataFrame(samples) - ds_out = Dataset.from_pandas(df) - fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}/{ds_name}/{split}/ds{shard_name}" + # df = pd.DataFrame(samples) + # ds_out = Dataset.from_pandas(df) + ds_out = Dataset.from_list(samples) + # fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}/{ds_name}/{split}/ds{shard_name}" + os.makedirs(os.path.dirname(fpath), exist_ok=True) fpath = f"{fpath}.parquet" ds_out.to_parquet(fpath) print(f"Saved to {fpath}") # Cleanup - del samples, df, ds_out, completions, messages, ctxs, questions + del samples, ds_out, completions, messages, ctxs, questions clear_gpu() @@ -369,6 +587,8 @@ if __name__ == "__main__": enable_prefix_caching=True, enable_chunked_prefill=True, max_model_len=MODEL_CTX_LEN.get(vllm_model), + max_num_batched_tokens=16384, + max_num_seqs=32, # avoid oom when getting logprobs ) print(f"{llm_kwargs=}") diff --git a/ds_config.json b/ds_config.json new file mode 100644 index 0000000..44a4344 --- /dev/null +++ b/ds_config.json @@ -0,0 +1,20 @@ +{ + "bf16": { + "enabled": "auto" + }, + "zero_optimization": { + "stage": 0, + "stage3_gather_16bit_weights_on_model_save": "no", + "offload_optimizer": { + "device": "auto" + }, + "offload_param": { + "device": "auto" + } + }, + "gradient_clipping": "auto", + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "gradient_accumulation_steps": "auto", + "steps_per_print": 2000000 +} \ No newline at end of file diff --git a/intx_sft.py b/intx_sft.py index 42460e3..728551d 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -17,7 +17,6 @@ from transformers import ( AutoConfig, set_seed, ) -from transformers.utils import is_liger_kernel_available from ctx_to_lora.configs import ( AggregatorArguments, @@ -45,6 +44,7 @@ from ctx_to_lora.metrics import ( ) from ctx_to_lora.model_loading import ( get_lora_config, + get_model, get_model_and_tokenizer, get_tokenizer, ) @@ -176,7 +176,7 @@ def main(): ############ Model setup if not ctx_args.from_pretrained_checkpoint: model_name = model_args.model_name_or_path - model, tokenizer = get_model_and_tokenizer( + base_model, tokenizer = get_model_and_tokenizer( **vars(model_args), train=True, requires_grad=False, # ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE, @@ -191,8 +191,8 @@ def main(): ctx_encoder_model_config = ctx_encoder_model_config.text_config ctx_tokenizer = get_tokenizer(ctx_name, train=True) else: - ctx_name = model.base_model.config.name_or_path - ctx_encoder_model_config = model.config + 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: @@ -200,7 +200,7 @@ def main(): logger.info("Using HyperLoRA") if not ctx_args.from_pretrained_checkpoint: hypernet_config = get_hypernet_config( - model, + base_model, ctx_encoder_model_config, hypernet_args, aggregator_args, @@ -215,13 +215,13 @@ def main(): " as the context encoder" ) model = ModulatedPretrainedModel( - model, + base_model, hypernet_config, ctx_encoder_args, ctx_args.use_kl_loss, ctx_args.use_sequence_packing, ) - training_args.gen_lora_l1_reg_coef = ctx_args.gen_lora_l1_reg_coef + else: logger.info( f"Loading from checkpoint: {ctx_args.from_pretrained_checkpoint}" @@ -237,7 +237,9 @@ def main(): 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.gen_lora_l1_reg_coef = ctx_args.gen_lora_l1_reg_coef + training_args.use_kl_loss = ctx_args.use_kl_loss if len([p for p in model.ctx_encoder.parameters() if p.requires_grad]): raise ValueError("ctx_encoder contains trainable parameters") @@ -268,6 +270,10 @@ def main(): tokenizer_kwargs = {"max_length": ctx_args.max_base_len} # not used ctx_tokenizer_kwargs = {"max_length": ctx_args.max_ctx_len} # not used for now + base_model = get_model( + **vars(model_args), train=False, requires_grad=False, peft_config=None + ) + _get_tokenized_dataset = partial( get_tokenized_dataset, max_qas_len=ctx_args.max_qas_len, @@ -284,6 +290,7 @@ def main(): add_negative_prompt=ctx_args.add_negative_prompt, use_kl_loss=ctx_args.use_kl_loss, set_format=None if ctx_args.use_sequence_packing else "pt", + base_model=base_model, # streaming=data_args.streaming, ) tokenized_ds = {"train": dict(), "validation": dict()} @@ -409,33 +416,24 @@ def main(): # 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 training_args.use_liger_kernel and is_liger_kernel_available(): - from liger_kernel.transformers import _apply_liger_kernel_to_instance - if isinstance(model, ModulatedPretrainedModel): - logger.info("Applying liger-kernel to 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("Applying liger-kernel to ctx_encoder_model") - ctx_base_model = model.ctx_encoder.base_model - _apply_liger_kernel_to_instance(model=ctx_base_model) - compile_linear(ctx_base_model) - # model.ctx_encoder.compile(fullgraph=True) - - elif isinstance(model, PeftModel): - logger.info("Applying liger-kernel to PeftModel") + if isinstance(model, ModulatedPretrainedModel): + if isinstance(model.base_model, PeftModel): + base_model = model.base_model.base_model + else: base_model = model.base_model - _apply_liger_kernel_to_instance(model=base_model.model) - # compile_linear(base_model) - base_model.compile(fullgraph=True, mode="max-autotune") - base_model.loss_function = torch.compile( - base_model.loss_function, fullgraph=True, mode="max-autotune" - ) + 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) + # model.ctx_encoder.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( diff --git a/scripts/short_ctx/gemma_qa_short_ctx_old_conf_distill.sh b/scripts/short_ctx/gemma_qa_short_ctx_old_conf_distill.sh new file mode 100644 index 0000000..7a75b34 --- /dev/null +++ b/scripts/short_ctx/gemma_qa_short_ctx_old_conf_distill.sh @@ -0,0 +1,43 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --nodes=1 +#SBATCH --partition=a3 +#SBATCH --gpus=4 +#SBATCH --output=slurm_logs/%x-%j.out +#SBATCH --error=slurm_logs/%x-%j.out + +port=$((10000 + ($SLURM_JOBID % 50000))) +echo "Using port: $port" + +# --gradient_accumulation_steps=8 --gradient_clipping=1.0 +uv run accelerate launch --main_process_port $port \ +--num_processes=4 --gpu_ids all intx_sft.py $1 \ +--model_name_or_path=google/gemma-2-2b-it \ +--num_train_epochs=5 \ +--per_device_train_batch_size=-1 \ +--gradient_accumulation_steps=16 \ +--per_device_eval_batch_size=64 \ +--target_modules=down_proj \ +--num_blocks=9 \ +--num_self_attn_per_block=0 \ +--n_latent_queries=208 \ +--num_pre_head_layers=1 \ +--lora_r=8 \ +--eval_strategy=no \ +--eval_steps=1000 \ +--save_steps=1000 \ +--learning_rate=4e-5 \ +--lora_dropout=0.0 \ +--neftune_noise_alpha=5 \ +--add_negative_prompt=False \ +--add_repeat_prompt=False \ +--use_sequence_packing=True \ +--max_qas_len=2048 \ +--max_qas_per_sample=1 \ +--max_packed_inp_len=4096 \ +--max_packed_ctx_len=4096 \ +--per_rank_gen=True \ +--per_layer_processing=True \ +--gen_lora_l1_reg_coef=0.1 \ +--logging_steps=50 \ +--use_kl_loss=True diff --git a/slurm_logs/.gitkeep b/slurm_logs/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index 2d25970..47c3b56 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -233,7 +233,7 @@ class TrainingArguments(TrainingArguments): default=100, ) use_liger_kernel: bool = field( - default=True, + default=False, ) remove_unused_columns: bool = field( default=False, diff --git a/src/ctx_to_lora/data/collator.py b/src/ctx_to_lora/data/collator.py index ae6af55..68ecee8 100644 --- a/src/ctx_to_lora/data/collator.py +++ b/src/ctx_to_lora/data/collator.py @@ -103,6 +103,14 @@ def train_collator(inp_list, tokenizer): if ctx_ids is not None: out["ctx_ids"] = ctx_ids out["ctx_attn_mask"] = ctx_attn_mask + if "logprobs_vals" in sample: + # for training with logits + logprobs_vals = [x.pop("logprobs_vals") for x in inp_list] + # logprobs_vals = tokenizer.pad( + # {"input_ids": logprobs_vals}, **padding_kwargs + # )["input_ids"] + out["logprobs_vals"] = logprobs_vals + out["logprobs_indices"] = [x.pop("logprobs_indices") for x in inp_list] return out diff --git a/src/ctx_to_lora/data/definitions.py b/src/ctx_to_lora/data/definitions.py index b3802b8..489e42f 100644 --- a/src/ctx_to_lora/data/definitions.py +++ b/src/ctx_to_lora/data/definitions.py @@ -10,6 +10,16 @@ TRANSFORMED_DATA_DIR = "data/processed_datasets" RAW_DATA_DIR = "data/raw_datasets/" SELF_GEN_DATA_DIR = f"{RAW_DATA_DIR}/self_gen/" +# see https://huggingface.co/datasets/YuxinJiang/LTE_train_data/viewer/default/train?row=97&views%5B%5D=train&sql_row=2 +# based on https://arxiv.org/pdf/2402.11905 +SELF_GEN_SYSTEM_MSG = ( + "### SYSTEM INSTRUCTION ###\n" + "You are an honest and helpful assistant.\n" + "You must use the information provided in the context for responding to the question.\n" + "**DO NOT** hallucinate or make up information.\n" + "### END OF SYSTEM INSTRUCTION ###" +) + REPEAT_PROMPTS = [ "Repeat the text.", @@ -124,23 +134,22 @@ DS_KWARGS = { split="train[60:]", ), ), - # "pwc_tiny": dict( - # train=dict(path="sggetao/PwC", split="train[900:2000]"), - # validation=dict(path="sggetao/PwC", split="train[:900]"), - # ), - "squad": dict( + "pwc_compact_tiny": dict( train=dict( - path="/home/rujikorn_sakana_ai/.cache/huggingface/hub/datasets--rajpurkar--squad", - split="train[900:]", - ), - validation=dict( - path="/home/rujikorn_sakana_ai/.cache/huggingface/hub/datasets--rajpurkar--squad", - split="train[:900]", - ), - test=dict( - path="/home/rujikorn_sakana_ai/.cache/huggingface/hub/datasets--rajpurkar--squad", - split="validation", + path="parquet", + data_files="data/raw_datasets/pwc_compact/train/ds.parquet", + # correspond to skipping to first 900 samples in the original dataset + split="train[60:200]", ), + ), + "pwc_tiny": dict( + train=dict(path="sggetao/PwC", split="train[900:2000]"), + validation=dict(path="sggetao/PwC", split="train[:900]"), + ), + "squad": dict( + train=dict(path="data/raw_datasets/squad", split="train[900:]"), + validation=dict(path="data/raw_datasets/squad", split="train[:900]"), + test=dict(path="data/raw_datasets/squad", split="validation"), # train=dict(path="rajpurkar/squad", split="train[900:]"), # validation=dict(path="rajpurkar/squad", split="train[:900]"), # test=dict(path="rajpurkar/squad", split="validation"), diff --git a/src/ctx_to_lora/data/packing.py b/src/ctx_to_lora/data/packing.py index 11aa612..aac38cc 100644 --- a/src/ctx_to_lora/data/packing.py +++ b/src/ctx_to_lora/data/packing.py @@ -95,8 +95,26 @@ def pack_data_points_FA( position_ids = np.empty(total_inp_len, dtype=np.long) labels = np.empty(total_inp_len, dtype=np.long) + has_logits = "logprobs_vals" in batch + + if has_logits: + sequences = zip( + batch["input_ids"], + batch["labels"], + batch["logprobs_vals"], + batch["logprobs_indices"], + ) + n_labels = sum(len(y) for x in batch["logprobs_vals"] for y in x) + k = len(batch["logprobs_vals"][0][0][0]) # assuming all have same k + logprobs_vals = np.empty((n_labels, k), dtype=np.float32) + logprobs_indices = np.empty((n_labels, k), dtype=np.int32) + logits_offset = 0 + else: + sequences = zip(batch["input_ids"], batch["labels"]) + offset = ctx_offset = 0 - for input_ids_b, labels_b in zip(batch["input_ids"], batch["labels"]): + for sample in sequences: + input_ids_b, labels_b = sample[:2] inp_start = offset if need_flatten: # compute position_ids for each sub-list in input_ids_b @@ -120,6 +138,17 @@ def pack_data_points_FA( labels[inp_start:inp_end] = labels_b offset += inp_len + if has_logits: + logprobs_vals_b, logprobs_indices_b = sample[2:] + logprobs_vals_b = concat_list(logprobs_vals_b) + logprobs_indices_b = concat_list(logprobs_indices_b) + logits_len = len(logprobs_vals_b) + logprobs_vals[logits_offset : logits_offset + logits_len] = logprobs_vals_b + logprobs_indices[logits_offset : logits_offset + logits_len] = ( + logprobs_indices_b + ) + logits_offset += logits_len + for ctx_ids_b in batch["ctx_ids"]: ctx_len = len(ctx_ids_b) ctx_start, ctx_end = ctx_offset, ctx_offset + ctx_len @@ -127,13 +156,17 @@ def pack_data_points_FA( ctx_position_ids[ctx_start:ctx_end] = np.arange(ctx_len, dtype=np.long) ctx_offset += ctx_len - return { + out = { "ctx_ids": ctx_ids, "ctx_position_ids": ctx_position_ids, "input_ids": input_ids, "position_ids": position_ids, "labels": labels, } + if has_logits: + out["logprobs_vals"] = logprobs_vals + out["logprobs_indices"] = logprobs_indices + return out def pack_batch( @@ -143,6 +176,8 @@ def pack_batch( max_packed_size: int = -1, ) -> dict[str, any]: need_flatten = check_is_iterable(batch["input_ids"][0][0]) + assert need_flatten + if need_flatten: n_queries = [len(x) for x in batch["input_ids"]] inp_lens = [[len(y) for y in x] for x in batch["input_ids"]] @@ -179,6 +214,10 @@ def pack_batch( "labels": [], "n_queries": [], } + has_logits = "logprobs_vals" in batch + if has_logits: + packed_batch["logprobs_vals"] = [] + packed_batch["logprobs_indices"] = [] packing_efficiency_ratios = [] ctx_packing_efficiency_ratios = [] @@ -192,6 +231,11 @@ def pack_batch( "input_ids": batch["input_ids"][start_idx:end_idx], "labels": batch["labels"][start_idx:end_idx], } + if has_logits: + group_items["logprobs_vals"] = batch["logprobs_vals"][start_idx:end_idx] + group_items["logprobs_indices"] = batch["logprobs_indices"][ + start_idx:end_idx + ] packed_item = pack_data_points_FA(group_items, need_flatten) packed_batch["ctx_ids"].append(packed_item["ctx_ids"]) packed_batch["ctx_position_ids"].append(packed_item["ctx_position_ids"]) @@ -199,6 +243,9 @@ def pack_batch( packed_batch["position_ids"].append(packed_item["position_ids"]) packed_batch["labels"].append(packed_item["labels"]) packed_batch["n_queries"].append(n_queries[start_idx:end_idx]) + if has_logits: + packed_batch["logprobs_vals"].append(packed_item["logprobs_vals"]) + packed_batch["logprobs_indices"].append(packed_item["logprobs_indices"]) if max_packed_inp_len > 0: inp_efficiency = len(packed_item["input_ids"]) / max_packed_inp_len diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index 79b9ab5..db30679 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -9,8 +9,10 @@ from os import path from typing import Any import datasets +import numpy as np +import torch from datasets import Dataset, is_caching_enabled, load_dataset -from transformers import PreTrainedTokenizerBase +from transformers import PreTrainedModel, PreTrainedTokenizerBase from ctx_to_lora.data.definitions import ( CLOSED_QA_INTX_TEMPLATES, @@ -20,6 +22,7 @@ from ctx_to_lora.data.definitions import ( RAW_DATA_DIR, REPEAT_PROMPTS, SELF_GEN_DATA_DIR, + SELF_GEN_SYSTEM_MSG, TRANSFORMED_DATA_DIR, ) from ctx_to_lora.data.packing import pack_batch @@ -130,7 +133,7 @@ def get_preprocessing_fn( "response": sample["answer"], } - elif ds_name == "pwc": + elif ds_name == "pwc" or ds_name == "pwc_tiny": # original pwc def f(sample): return { @@ -309,6 +312,22 @@ def get_preprocessing_fn( f = maybe_convert_to_list(f) + if "self_gen" not in ds_name: + + def strip_response(f): + def g(sample): + sample = f(sample) + if "responses" in sample and bool(sample["responses"]): + sample["responses"] = [ + r.strip() if isinstance(r, str) else r + for r in sample["responses"] + ] + return sample + + return g + + f = strip_response(f) + return f @@ -551,7 +570,19 @@ def load_and_process_dataset( cols_to_remove = [ col for col in ds.column_names - if col not in ["context", "prompts", "responses", "qas", "variation"] + if col + not in [ + "context", + "prompts", + "responses", + "qas", + "variation", + "logprobs_vals", + "logprobs_indices", + "input_ids", + "ctx_ids", + "response_start_end", + ] ] is_eval = split != "train" ds = ds.map( @@ -613,9 +644,9 @@ def get_tokenized_dataset( repeat_prob: float = 0.0, # only used when add_repeat_prompt is True set_format: str | None = None, streaming: bool = False, + base_model: PreTrainedModel | None = None, ) -> dict[str, Any]: - # TODO (distillation, Tan): make this works with pre-computed logits - assert not use_kl_loss, "KL loss is deprecated" + # assert not use_kl_loss, "KL loss is deprecated" # TO BE REMOVED if add_repeat_prompt: assert repeat_prob > 0, f"add_repeat_prompt is set but repeat_prob = 0" logger.info(f"Loading dataset {ds_name} with split {split}...") @@ -673,6 +704,7 @@ def get_tokenized_dataset( ds=ds, tokenizer=tokenizer, ctx_tokenizer=ctx_tokenizer, + model=base_model, num_proc=num_proc, **tokenize_kwargs, ) @@ -696,6 +728,7 @@ def construct_and_tokenize_ctx_qa( ctx_tokenizer_kwargs, add_ctx_to_chat, use_kl_loss, + model, need_ctx_ids, is_pretrain, is_paraphrased, @@ -740,74 +773,104 @@ def construct_and_tokenize_ctx_qa( # for sft + chat_model, we need to convert the dataset to chat format # add "data" field - ds = ds.map( - convert_ctx_prompt_response_to_messages, - fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat}, - num_proc=num_proc, - remove_columns=[col for col in ds.column_names if col != "context"], - ) - # add input_ids, attention_mask, labels to "data" - os.environ["TOKENIZERS_PARALLELISM"] = "true" - logging.debug("Tokenizing inputs") - tokenized_ds = ds.map( - get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer), - batched=True, - batch_size=100_000, - # # num_proc=num_proc, - ) + if "input_ids" in ds.column_names and "response_start_end" in ds.column_names: + # already tokenized dataset (e.g., self-gen qa data) + tokenized_ds = ds.map( + get_labels_from_input_ids, + num_proc=num_proc, + ) + else: + ds = ds.map( + convert_ctx_prompt_response_to_messages, + fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat}, + num_proc=num_proc, + # remove_columns=[col for col in ds.column_names if col != "context"], + ) + # add input_ids, attention_mask, labels to "data" + os.environ["TOKENIZERS_PARALLELISM"] = "true" + logging.debug("Tokenizing inputs") + tokenized_ds = ds.map( + get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer), + batched=True, + batch_size=100_000, + # # num_proc=num_proc, + ) - if "train" in split: - # base model cant handle these samples naturally - # should skip - # tokenized_ds = tokenized_ds.filter( - # len_filter, - # fn_kwargs={ - # "max_length": base_model_max_len, - # "keys": ["input_ids", "labels"], - # }, - # # batched=False, - # # batch_size=0, - # num_proc=num_proc, - # ) - tokenized_ds = tokenized_ds.map( - len_filter_data, - fn_kwargs={ - "max_length": base_model_max_len, - "keys": ["input_ids"], - }, - num_proc=num_proc, - ) + # if "train" in split: + # # base model cant handle these samples naturally + # # should skip + # # tokenized_ds = tokenized_ds.filter( + # # len_filter, + # # fn_kwargs={ + # # "max_length": base_model_max_len, + # # "keys": ["input_ids", "labels"], + # # }, + # # # batched=False, + # # # batch_size=0, + # # num_proc=num_proc, + # # ) + # tokenized_ds = tokenized_ds.map( + # len_filter_data, + # fn_kwargs={ + # "max_length": base_model_max_len, + # "keys": ["input_ids"], + # }, + # num_proc=num_proc, + # ) - # for use_kl_loss, we need "chat_ids" and "chat_attn_mask" - # TODO (distillation): implement + logger.info(tokenized_ds) + + if use_kl_loss and ("train" in split): + if "logprobs_vals" not in ds.column_names: + assert model is not None, ( + f"Base model is required for computing the " + "logits when `use_kl_loss` is True and logprobs are not included." + ) + # add "data_logp" + tokenized_ds = tokenized_ds.map( + convert_ctx_prompt_response_to_messages, + fn_kwargs={"add_ctx_to_chat": True, "for_logp": True}, + num_proc=num_proc, + ) + # add logits info to "data_logp" field + logger.info("computing logprobs") + tokenized_ds = tokenized_ds.map( + get_sft_prompt_formatting_fn( + TRAINING_TASK.COMPLETION, tokenizer, use_kl_loss, model + ), + batched=True, + batch_size=100_000, + ) + else: + tokenized_ds = tokenized_ds.map(add_ctx_attn_mask, num_proc=num_proc) + + keep_cols = [ + "input_ids", + "attention_mask", + "labels", + "ctx_ids", + "ctx_attn_mask", + "context", + "data", + "data_logp", + ] if use_kl_loss: - raise NotImplementedError("KL loss not implemented") - - # old code from before (might not work anymore) - # e.g,. the next line likely replace the tokenized text from - # earlier `convert_ctx_prompt_response_to_messages` call - - # tokenized_ds = tokenized_ds.map( - # convert_ctx_prompt_response_to_messages, - # fn_kwargs={"add_ctx_to_chat": True}, - # remove_columns=["messages"], - # ) - # tokenized_ds = tokenized_ds.map( - # get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer), - # remove_columns=["chat"], - # ) + keep_cols += ["logprobs_vals", "logprobs_indices"] + tokenized_ds = tokenized_ds.remove_columns( + [col for col in tokenized_ds.column_names if col not in keep_cols], + ) if need_ctx_ids: - # tokenize the ctx_text to get ctx_ids and ctx_attn_mask - os.environ["TOKENIZERS_PARALLELISM"] = "true" - logging.debug("Tokenizing context") - tokenized_ds = tokenized_ds.map( - tokenize_ctx_text, - fn_kwargs={"tokenizer": ctx_tokenizer}, - batched=True, - batch_size=100_000, - # # num_proc=num_proc, - ) + if "ctx_ids" not in tokenized_ds.column_names: + # tokenize the ctx_text to get ctx_ids and ctx_attn_mask + os.environ["TOKENIZERS_PARALLELISM"] = "true" + logging.debug("Tokenizing context") + tokenized_ds = tokenized_ds.map( + tokenize_ctx_text, + fn_kwargs={"tokenizer": ctx_tokenizer}, + batched=True, + batch_size=100_000, + ) if "train" in split: # TODO: do something if ctx length is longer than the ctx model length # e.g., drop or split for multi-lora training @@ -823,8 +886,13 @@ def construct_and_tokenize_ctx_qa( tokenized_ds = tokenized_ds.map( unpack_data, num_proc=num_proc, - remove_columns=["data", "context"], + remove_columns=[ + col + for col in tokenized_ds.column_names + if col in ["data", "data_logp", "context"] + ], ) + print(tokenized_ds) logging.debug(f"Split too long QAs with max length {max_qas_len}") tokenized_ds = tokenized_ds.map( @@ -900,9 +968,57 @@ def build_intx_pretrain(sample: dict[str, str]): return {"text": sample["context"] + sample["qas"]} +def add_ctx_attn_mask(sample: dict[str, Any]) -> dict[str, Any]: + """ + Add context attention mask to the sample. + + Args: + sample: A dictionary containing 'ctx_ids' + + Returns: + The sample with 'ctx_attn_mask' added + """ + if "ctx_ids" in sample: + sample["ctx_attn_mask"] = [1] * len(sample["ctx_ids"]) + return sample + + +def get_labels_from_input_ids(sample: dict[str, Any]) -> dict[str, Any]: + """ + Extract labels from input_ids and response_start. + + Args: + sample: A dictionary containing 'input_ids' and 'response_start' + + Returns: + A dictionary with 'labels' field added + """ + labels = [] + attention_mask = [] + for input_ids_i, (start_i, end_i) in zip( + sample["input_ids"], sample["response_start_end"] + ): + len_input_ids = len(input_ids_i) + # pad labels with -100 + pad_len_left = start_i + pad_len_right = len_input_ids - end_i + labels.append( + [IGNORE_INDEX] * pad_len_left + + input_ids_i[start_i:] + + [IGNORE_INDEX] * pad_len_right + ) + attention_mask.append([1] * len(input_ids_i)) + + sample["labels"] = labels + sample["attention_mask"] = attention_mask + return sample + + def get_sft_prompt_formatting_fn( sft_mode: TRAINING_TASK, tokenizer: PreTrainedTokenizerBase, + for_logp: bool = False, + model: PreTrainedModel | None = None, ) -> Callable[[dict[str, Any]], dict[str, Any]]: """ Get a function that formats examples for supervised fine-tuning. @@ -929,10 +1045,13 @@ def get_sft_prompt_formatting_fn( "Training with pre-training data is not supported yet." ) + data_field = "data_logp" if for_logp else "data" + + @torch.inference_mode() def f_intx(samples): # flatten all the messages into a list # tokenize, the pack back correctly - messages_list = [x for x in samples["data"]] + messages_list = [x for x in samples[data_field]] n_queries = [len(x) for x in messages_list] messages = concat_list(messages_list) @@ -940,31 +1059,76 @@ def get_sft_prompt_formatting_fn( tokens = tokenizer.apply_chat_template( messages, tokenize=True, - add_special_token=False, + add_special_tokens=False, + padding=False, truncation=False, add_generation_prompt=False, return_assistant_tokens_mask=True, return_dict=True, ) labels = [] + logprobs_vals = [] + logprobs_indices = [] for tok_ids, masks in zip(tokens["input_ids"], tokens["assistant_masks"]): o = [id_ if mask else IGNORE_INDEX for id_, mask in zip(tok_ids, masks)] labels.append(o) + if for_logp: + # TODO: batch compute? + # save only the logits for the assistant tokens + assert sum(masks) > 0, "No assistant tokens found in the masks" + + # + 1 to also get the logits of the token before the first label token + num_logits_to_keep = int(len(masks) - np.argmax(masks) + 1) + + # [n_labels + 1] + logits = model( + torch.tensor([tok_ids], device=model.device), + logits_to_keep=num_logits_to_keep, + ).logits.squeeze(0) + + # include the logits of the token before the first label token + pre_label_idx = np.argmax(masks) - 1 + masks[pre_label_idx] = True + logits = logits[ + torch.tensor( + masks[pre_label_idx:], + device=logits.device, + dtype=torch.bool, + ) + ] + + logits = logits[..., :-1, :] # remove the last token (eos) + logprobs = torch.nn.functional.log_softmax(logits.float(), dim=-1) + # ignore eos + vals, idx = torch.topk(logprobs, k=16, dim=-1) + logprobs_vals.append(vals.to(torch.float16)) + logprobs_indices.append(idx.to(torch.int32)) + del tokens["assistant_masks"] tokens["labels"] = labels per_ctx_tokens = [] i = 0 for n in n_queries: - per_ctx_tokens.append( - { + if for_logp: + per_ctx_data = { + "logprobs_vals": logprobs_vals[i : i + n], + "logprobs_indices": logprobs_indices[i : i + n], + } + # per_ctx_data = { + # "logprobs_vals": logprobs_vals[i : i + n], + # "logprobs_indices": logprobs_indices[i : i + n], + # } + else: + per_ctx_data = { "input_ids": tokens["input_ids"][i : i + n], "attention_mask": tokens["attention_mask"][i : i + n], "labels": tokens["labels"][i : i + n], } - ) + + per_ctx_tokens.append(per_ctx_data) i += n - return dict(data=per_ctx_tokens) + return {data_field: per_ctx_tokens} # return f if not apply_chat_template_fn is not None else f_intx return f_intx @@ -973,6 +1137,7 @@ def get_sft_prompt_formatting_fn( def convert_ctx_prompt_response_to_messages( example: dict[str, Any], add_ctx_to_chat: bool, + for_logp: bool = False, ) -> dict[str, Any]: """ Convert context/prompt/response format to chat messages format. @@ -998,6 +1163,8 @@ def convert_ctx_prompt_response_to_messages( system_msg = "" if "system_message" in example: system_msg = example["system_message"].strip() + if for_logp and not system_msg: + system_msg = SELF_GEN_SYSTEM_MSG data = [] for prompt, response in zip(example[prompt_field], example[res_field]): @@ -1009,11 +1176,12 @@ def convert_ctx_prompt_response_to_messages( [ {"role": "system", "content": system_msg.strip()}, {"role": "user", "content": user_msg.strip()}, - {"role": "assistant", "content": response.strip()}, + {"role": "assistant", "content": response}, ] ) + out_field = "data_logp" if for_logp else "data" - return dict(data=data) + return {out_field: data} def split_too_long_qas( @@ -1031,6 +1199,8 @@ def split_too_long_qas( labels = samples["labels"] ctx_ids = samples["ctx_ids"] ctx_attn_mask = samples["ctx_attn_mask"] + target_logprobs_vals = samples.get("logprobs_vals", None) + target_logprobs_indices = samples.get("logprobs_indices", None) # Pre-calculate total lengths to check if any splitting is needed total_lengths = [sum(len(x) for x in seq) for seq in input_ids] @@ -1048,14 +1218,28 @@ def split_too_long_qas( out = {k: list() for k in samples} longest_new_qas_len = 0 n_skip = 0 + has_target_logprobs = ( + target_logprobs_vals is not None and target_logprobs_indices is not None + ) # Helper function to add a batch efficiently - def add_batch(inp_ids_batch, attn_batch, labels_batch, ctx_id, ctx_attn): + def add_batch( + inp_ids_batch, + attn_batch, + labels_batch, + ctx_id, + ctx_attn, + target_vals_batch=None, + target_indices_batch=None, + ): out["input_ids"].append(inp_ids_batch) out["attention_mask"].append(attn_batch) out["labels"].append(labels_batch) out["ctx_ids"].append(ctx_id) out["ctx_attn_mask"].append(ctx_attn) + if has_target_logprobs: + out["logprobs_vals"].append(target_vals_batch) + out["logprobs_indices"].append(target_indices_batch) for i, tot_inp_len in enumerate(total_lengths): if (max_qas_len < 0 or tot_inp_len <= max_qas_len) and ( @@ -1074,10 +1258,26 @@ def split_too_long_qas( new_input_ids = [] new_attn_mask = [] new_labels = [] + new_target_vals = [] if has_target_logprobs else None + new_target_indices = [] if has_target_logprobs else None + + sequences = zip(input_ids[i], attention_mask[i], labels[i]) + if has_target_logprobs: + sequences = zip( + input_ids[i], + attention_mask[i], + labels[i], + target_logprobs_vals[i], + target_logprobs_indices[i], + ) + + for seq_data in sequences: + if has_target_logprobs: + inp_ids, attn_mask, label, target_vals, target_indices = seq_data + else: + inp_ids, attn_mask, label = seq_data + target_vals, target_indices = None, None - for inp_ids, attn_mask, label in zip( - input_ids[i], attention_mask[i], labels[i] - ): inp_len = len(inp_ids) if inp_len > max_qas_len: # Skip individual sequences that are too long @@ -1095,6 +1295,9 @@ def split_too_long_qas( new_input_ids.append(inp_ids) new_attn_mask.append(attn_mask) new_labels.append(label) + if has_target_logprobs: + new_target_vals.append(target_vals) + new_target_indices.append(target_indices) else: # Current batch is full, save it and start new batch # logger.debug( @@ -1107,6 +1310,8 @@ def split_too_long_qas( new_labels, current_ctx_id, current_ctx_attn, + new_target_vals, + new_target_indices, ) longest_new_qas_len = max(longest_new_qas_len, new_qas_len) @@ -1115,6 +1320,9 @@ def split_too_long_qas( new_input_ids = [inp_ids] new_attn_mask = [attn_mask] new_labels = [label] + if has_target_logprobs: + new_target_vals = [target_vals] + new_target_indices = [target_indices] # Add final batch if not empty if new_input_ids: @@ -1124,6 +1332,8 @@ def split_too_long_qas( new_labels, current_ctx_id, current_ctx_attn, + new_target_vals, + new_target_indices, ) longest_new_qas_len = max(longest_new_qas_len, new_qas_len) @@ -1138,7 +1348,12 @@ def split_too_long_qas( def unpack_data(sample): - return {k: sample["data"][k] for k in sample["data"]} + if "data" not in sample: + return sample + data_keys = ["data"] + if "data_logp" in sample: + data_keys.append("data_logp") + return {k: sample[dk][k] for dk in data_keys for k in sample[dk]} def unpack_data_eval(samples): @@ -1317,6 +1532,7 @@ def tokenize_pretrain( tokens = tokenizer( samples["text"], add_special_tokens=True, + padding=False, truncation=False, **(tokenizer_kwargs or {}), ) @@ -1438,6 +1654,72 @@ def pack( return ds +# def add_logits(model, ds): +# # compute logits for prepared tokenized ds +# # must handle packed and flattened dataset +# # also cache! +# # let's do this row by row +# # slower for non-packed ds +# input = torch.tensor([ds['input_ids']], device=model.device) +# labels = torch.tensor([ds['labels']], device=model.device) +# with torch.no_grad(): +# logits = model(input).logits + +# top_logits, top_logits_positions = torch.topk(logits, args.max_logits_to_store, dim=-1) + +# # Convert logits to specified precision +# dtype_map = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16} +# target_dtype = dtype_map[args.logit_precision] +# top_logits = top_logits.to(target_dtype) + +# tokens_to_mask = tok.apply_chat_template( +# [messages[0]], # Just the user message +# tokenize=True, +# add_special_tokens=False, +# truncation=False, +# padding=False, +# add_generation_prompt=True, +# return_dict=True, +# ) +# number_of_tokens_to_mask = len(tokens_to_mask['input_ids']) + +# # Create mask: False for prompt tokens, True for response tokens +# responses_mask = torch.zeros(len(all_tokens['input_ids']), dtype=torch.bool) +# responses_mask[number_of_tokens_to_mask:] = True + +# responses_masks.append(responses_mask) +# all_logits_values.append(top_logits.cpu()) +# all_logits_positions.append(top_logits_positions.cpu()) + +# logits = model(input).logits + +# top_logits, top_logits_positions = torch.topk(logits, args.max_logits_to_store, dim=-1) + +# # Convert logits to specified precision +# dtype_map = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16} +# target_dtype = dtype_map[args.logit_precision] +# top_logits = top_logits.to(target_dtype) + +# tokens_to_mask = tok.apply_chat_template( +# [messages[0]], # Just the user message +# tokenize=True, +# add_special_tokens=False, +# truncation=False, +# padding=False, +# add_generation_prompt=True, +# return_dict=True, +# ) +# number_of_tokens_to_mask = len(tokens_to_mask['input_ids']) + +# # Create mask: False for prompt tokens, True for response tokens +# responses_mask = torch.zeros(len(all_tokens['input_ids']), dtype=torch.bool) +# responses_mask[number_of_tokens_to_mask:] = True + +# responses_masks.append(responses_mask) +# all_logits_values.append(top_logits.cpu()) +# all_logits_positions.append(top_logits_positions.cpu()) + + if __name__ == "__main__": from transformers import AutoTokenizer diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py index fa0a8b5..b4f1476 100644 --- a/src/ctx_to_lora/eval_utils.py +++ b/src/ctx_to_lora/eval_utils.py @@ -14,15 +14,12 @@ import pandas as pd import torch import yaml from datasets import disable_caching -from peft import PeftModel from transformers import ( - PreTrainedModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, Trainer, set_seed, ) -from transformers.utils import is_liger_kernel_available from ctx_to_lora.data.collator import eval_collator, generation_collator from ctx_to_lora.data.definitions import ( @@ -693,24 +690,24 @@ def evaluate( use_flash_attn=True, ) - if is_liger_kernel_available(): - from liger_kernel.transformers import _apply_liger_kernel_to_instance + # if is_liger_kernel_available(): + # from liger_kernel.transformers import _apply_liger_kernel_to_instance - if isinstance(model, ModulatedPretrainedModel): - print("Applying liger-kernel to ModulatedPretrainedModel") - if isinstance(model.base_model, PeftModel): - _apply_liger_kernel_to_instance(model=model.base_model.base_model.model) - else: - _apply_liger_kernel_to_instance(model=model.base_model.model) - if ctx_name is not None: - print("Applying liger-kernel to ctx_encoder_model") - _apply_liger_kernel_to_instance(model=model.ctx_encoder.base_model) - elif isinstance(model, PeftModel): - print("Applying liger-kernel to PeftModel") - _apply_liger_kernel_to_instance(model=model.base_model.model) - elif isinstance(model, PreTrainedModel): - print("Applying liger-kernel to PretrainedModel") - _apply_liger_kernel_to_instance(model=model.model) + # if isinstance(model, ModulatedPretrainedModel): + # print("Applying liger-kernel to ModulatedPretrainedModel") + # if isinstance(model.base_model, PeftModel): + # _apply_liger_kernel_to_instance(model=model.base_model.base_model.model) + # else: + # _apply_liger_kernel_to_instance(model=model.base_model.model) + # if ctx_name is not None: + # print("Applying liger-kernel to ctx_encoder_model") + # _apply_liger_kernel_to_instance(model=model.ctx_encoder.base_model) + # elif isinstance(model, PeftModel): + # print("Applying liger-kernel to PeftModel") + # _apply_liger_kernel_to_instance(model=model.base_model.model) + # elif isinstance(model, PreTrainedModel): + # print("Applying liger-kernel to PretrainedModel") + # _apply_liger_kernel_to_instance(model=model.model) tokenizer = get_tokenizer(args.model_name_or_path, train=False) if tokenizer.pad_token_id is None: diff --git a/src/ctx_to_lora/modeling/lora_layer.py b/src/ctx_to_lora/modeling/lora_layer.py index 2a44852..94b7957 100644 --- a/src/ctx_to_lora/modeling/lora_layer.py +++ b/src/ctx_to_lora/modeling/lora_layer.py @@ -29,11 +29,12 @@ def lora_forward( B = B.repeat_interleave(n_qs, dim=0, output_size=tot_q) base_out = torch.nn.Linear.forward(self, x, *args, **kwargs) + x = x.to(A.dtype) delta_x = F.dropout(x, p=lora_dropout_p, training=self.training) delta_x = einsum(A, delta_x, "tot_q r d_in, tot_q s_len d_in -> tot_q s_len r") delta_x = einsum(B, delta_x, "tot_q d_out r, tot_q s_len r -> tot_q s_len d_out") delta_x = delta_x * scaling - return base_out + delta_x + return (base_out + delta_x).to(base_out.dtype) def lora_forward_packed( @@ -52,7 +53,7 @@ def lora_forward_packed( ) -> Float[Tensor, "1 tot_len d_out"]: # bs of x should be 1 in this case base_out = torch.nn.Linear.forward(self, x, *args, **kwargs) - + x = x.to(A.dtype) delta_x = F.dropout(x, p=lora_dropout_p, training=self.training) # # [tot_seq_len, r, d_in] # repeated_A = A.repeat_interleave(seq_lens, dim=0, output_size=seq_lens.sum()) @@ -73,7 +74,7 @@ def lora_forward_packed( ) delta_x = delta_x * scaling - return base_out + delta_x + return (base_out + delta_x).to(base_out.dtype) def apply_lora_to_layers( diff --git a/src/ctx_to_lora/trainer.py b/src/ctx_to_lora/trainer.py index ef1c382..b2593fb 100644 --- a/src/ctx_to_lora/trainer.py +++ b/src/ctx_to_lora/trainer.py @@ -1,5 +1,6 @@ import logging +import torch from torch import nn from transformers import Trainer from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES @@ -12,8 +13,126 @@ from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel logger = logging.getLogger() -# TODO (distillation, Shin): implement -class DistillationTrainer(Trainer): ... +def add_pre_label_token(labels, is_packed): + """ + Add a token before the first label token. + 1 is used as a placeholder for the first label token. + We compute the loss against the teacher logits anyway. + """ + if is_packed: + non_label_mask = torch.where(labels == -100, True, False) + diff = torch.diff(labels, append=torch.tensor([[-100]], device=labels.device)) + pre_label_pos = torch.where(non_label_mask * (diff != 0)) + labels[pre_label_pos] = 1 + else: + bsz = labels.shape[0] + first_label_pos = torch.argmax((labels != -100).float(), axis=-1) + labels[torch.arange(bsz), first_label_pos - 1] = 1 + return labels + + +# TODO: test +class DistillationTrainer(Trainer): + def __init__(self, *args, **kwargs): + self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0) + self.distillation_temperature = kwargs.pop("distillation_temperature", 1.0) + super().__init__(*args, **kwargs) + + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): + if self.model_accepts_loss_kwargs: + loss_kwargs = {} + if num_items_in_batch is not None: + loss_kwargs["num_items_in_batch"] = num_items_in_batch + inputs = {**inputs, **loss_kwargs} + + labels = inputs.pop("labels", None) # remove labels if present + # NOTE: have to account that we save +1 # token before the first label token + is_packed = "position_ids" in inputs + # labels = add_pre_label_token(labels, is_packed) + label_pos = torch.where(labels != -100) + # TODO: add efficient logits indexing? + # https://github.com/huggingface/trl/pull/2773 + # + use `num_logits_to_keep` to limit the number of logits + outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True) + target_logp = inputs.pop("logprobs_vals") + indices = inputs.pop("logprobs_indices") + + # packed + if isinstance(target_logp, torch.Tensor): + target_logp = target_logp.squeeze(0) + indices = indices.squeeze(0) + # non-packed max_qas_per_sample=1 + elif isinstance(target_logp[0], torch.Tensor): + target_logp = torch.cat([x[0] for x in target_logp], dim=0) + indices = torch.cat([x[0] for x in indices], dim=0) + # non-packed max_qas_per_sample>1 + else: + # convert to tensor + raise NotImplementedError( + "Non-packed data with max_qas_per_sample > 1 is not supported yet." + ) + + assert label_pos[0].shape[0] == target_logp.shape[0], ( + "Label positions and target log probabilities should have the same # tokens." + f"Got : {label_pos[0].shape[0]=} and {target_logp.shape[0]=}" + ) + + # TODO: per-ctx average loss? + ##### KL loss + outputs_logits = outputs.logits[label_pos[0], label_pos[1] - 1] # shift back 1 + + teacher_logp = torch.full_like(outputs_logits, -torch.inf) + teacher_logp.scatter_(1, indices, target_logp) + reduction = "batchmean" if num_items_in_batch is None else "sum" + + # loss = torch.nn.functional.kl_div( + # torch.nn.functional.log_softmax( + # outputs_logits / self.distillation_temperature, dim=-1 + # ), + # torch.nn.functional.softmax( + # teacher_logits / self.distillation_temperature, dim=-1 + # ), + # reduction=reduction, + # log_target=False, + # ) * (self.distillation_temperature**2) + # loss = loss / num_items_in_batch if num_items_in_batch is not None else loss + + # p = torch.nn.functional.softmax( + # teacher_logits / self.distillation_temperature, dim=-1 + # ) + p = teacher_logp.exp() + logq = nn.functional.log_softmax(outputs_logits, dim=-1) + + loss = -torch.sum(p * logq, dim=-1) * (self.distillation_temperature**2) + if reduction == "batchmean": + loss = loss.mean() + elif reduction == "sum" and num_items_in_batch is not None: + loss = loss.sum() / num_items_in_batch + ##### + + ##### unpack gen lora dict and compute regularization loss + l1_norm = 0 + for module, lora in gen_loras.items(): + l1_norm += lora["A"].abs().mean() + lora["B"].abs().mean() + # rough estimate of the losses (we only log the values from one step) + if (self.state.global_step == 1 and self.args.logging_first_step) or ( + self.args.logging_strategy == IntervalStrategy.STEPS + and self.state.global_step % self.state.logging_steps == 0 + ): + self.log({"ce_loss": loss.item(), "gen_lora_l1_norm": l1_norm.item()}) + loss += self.gen_lora_l1_reg_coef * l1_norm + ##### + + if ( + self.args.average_tokens_across_devices + and self.model_accepts_loss_kwargs + and num_items_in_batch is not None + ): + loss *= self.accelerator.num_processes + + return (loss, outputs) if return_outputs else loss class ModulatedModelTrainer(Trainer): @@ -40,6 +159,9 @@ class ModulatedModelTrainer(Trainer): if num_items_in_batch is not None: loss_kwargs["num_items_in_batch"] = num_items_in_batch inputs = {**inputs, **loss_kwargs} + # TODO: add efficient logits indexing? + # https://github.com/huggingface/trl/pull/2773 + # + use `num_logits_to_keep` to limit the number of logits outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True) # Save past state if it exists # TODO: this needs to be fixed and made cleaner later. @@ -86,7 +208,11 @@ class ModulatedModelTrainer(Trainer): loss += self.gen_lora_l1_reg_coef * l1_norm ##### - if self.args.average_tokens_across_devices and self.model_accepts_loss_kwargs: + if ( + self.args.average_tokens_across_devices + and self.model_accepts_loss_kwargs + and num_items_in_batch is not None + ): loss *= self.accelerator.num_processes return (loss, outputs) if return_outputs else loss @@ -122,8 +248,6 @@ def train_model( checkpoint = training_args.resume_from_checkpoint logger.info(f"Resuming from the checkpoint: {checkpoint}") - is_modulated_model = isinstance(model, ModulatedPretrainedModel) - trainer_cls = Trainer if not is_modulated_model else ModulatedModelTrainer trainer_kwargs = dict( model=model, args=training_args, @@ -132,10 +256,20 @@ def train_model( data_collator=train_collator, compute_metrics=compute_metrics, ) + + is_modulated_model = isinstance(model, ModulatedPretrainedModel) + trainer_cls = Trainer if is_modulated_model: - logger.info("Training with modulated model. Using CustomTrainer.") + logger.info("Training with modulated model. Using ModulatedModelTrainer.") + trainer_cls = ModulatedModelTrainer trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef del training_args.gen_lora_l1_reg_coef + + if training_args.use_kl_loss: + logger.info("Training with distillation loss. Using DistillationTrainer.") + trainer_cls = DistillationTrainer + del training_args.use_kl_loss + if training_args.auto_find_batch_size: # set the batch size to some high number # which will be lowered by the Trainer