From 650395874faed620e3caf2a1f3730b81133e421a Mon Sep 17 00:00:00 2001 From: 51616 Date: Sun, 7 Sep 2025 17:16:20 +0900 Subject: [PATCH] toy exp w/ self-gen + use double bos for eval (following vllm 0.8.4) --- .github/instructions/.instructions.md | 17 ++++--- README.md | 11 ++++- .../ctx_magic_number_32_256_self_gen.yaml | 18 ++++++++ data/self_generate_qa.py | 44 +++++++++++++++---- scripts/toy_exp.sh | 8 +++- src/ctx_to_lora/configs.py | 2 +- src/ctx_to_lora/data/processing.py | 7 ++- src/ctx_to_lora/trainer.py | 3 ++ watcher.py | 30 +++++++------ 9 files changed, 109 insertions(+), 31 deletions(-) create mode 100644 configs/toy_exp/ctx_magic_number_32_256_self_gen.yaml diff --git a/.github/instructions/.instructions.md b/.github/instructions/.instructions.md index ab710b2..3c7c3c2 100644 --- a/.github/instructions/.instructions.md +++ b/.github/instructions/.instructions.md @@ -1,11 +1,18 @@ --- applyTo: '**' --- -When asked for a change or implementation, only work on the specific task I have given you with the most concise and elegant solution that changes as little code as possible. +# Implementation +- When asked for a change or implementation, only work on the specific task I have given you with the most concise and elegant solution that changes as little code as possible. +- If the requested feature has unintended side effect, let the user know before implementing. -When reviewing, first check the code diff for potential bugs and inconsistencies. -After that, check if the change would cause any bugs somewhere else in the codebase. +# Review +- When reviewing, first check the code diff for potential bugs and inconsistencies. After that, check if the change would cause any bugs somewhere else in the codebase. -You should give detailed, professional, and nicely formatted explanation. +# Code Style +- Follow best practices to keep code clean and succinct. + +# General +- You should give detailed, professional, and nicely formatted explanation. +- Always think slowly and carefully before doing anything. Be objective. +- If instruction is unclear or ambiguous, ask the user for clarification. -Always think slowly and carefully before doing anything. Be objective. diff --git a/README.md b/README.md index 1833b1c..ec22059 100644 --- a/README.md +++ b/README.md @@ -283,4 +283,13 @@ srv # copy-paste the relative path from the root # e.g., http://localhost:8001/train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/eval-results-105000/comparator.json -``` \ No newline at end of file +``` + +## :warning: Known Issues +For the ICLR version, we use two `` tokens for `gemma-2-2b-it`. +This is because of a bug from `vllm` which we didn't realize until very close to the deadline (7 Sep 2025) (see https://github.com/vllm-project/vllm/pull/16081). +This bug causes the self-generated data to have two `` tokens (only the `input_ids` not `ctx_ids`) + +NOTE: previously we were using `vllm==0.8.5` which does not have the double `` problem. +However it has another bug that randomly causes some detokenization error (see https://github.com/vllm-project/vllm/issues/17448) +Therefore we, not knowing about the double `` bug, degraded `vllm` to `0.8.4`. diff --git a/configs/toy_exp/ctx_magic_number_32_256_self_gen.yaml b/configs/toy_exp/ctx_magic_number_32_256_self_gen.yaml new file mode 100644 index 0000000..eef0b8b --- /dev/null +++ b/configs/toy_exp/ctx_magic_number_32_256_self_gen.yaml @@ -0,0 +1,18 @@ +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +use_kl_loss: true + +# data +train_ds_names: +- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.0_summ_prob_0.0_struct_prob_0.0/ctx_magic_number_32_128 +- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.0_summ_prob_0.0_struct_prob_0.0/ctx_magic_number_128_256 + +val_ds_names: +- ctx_magic_number_32_128 +- ctx_magic_number_128_256 + + diff --git a/data/self_generate_qa.py b/data/self_generate_qa.py index 93f4fe7..d6f8756 100644 --- a/data/self_generate_qa.py +++ b/data/self_generate_qa.py @@ -53,7 +53,8 @@ SELF_QA_INTX = ( PRE_CTX = "# Provided Information\n" -PROMPT_TEMPLATE = PRE_CTX + "{context}\n\n---\n\n" + SELF_QA_INTX + "{question}" +QA_PROMPT_TEMPLATE = PRE_CTX + "{context}\n\n---\n\n" + SELF_QA_INTX + "{question}" +PROMPT_TEMPLATE = "{context}\n\n{question}" MODEL_CTX_LEN = { "google/gemma-2-27b-it": 8192, @@ -85,8 +86,10 @@ def truncate_middle_if_too_long( return input_ids -def get_prompt(context: str, q: str) -> str: - return PROMPT_TEMPLATE.format(context=context, question=q) +# TODO: self-gen for magic ctx num +def get_prompt(context: str, q: str, remove_qa_template: bool) -> str: + prompt = QA_PROMPT_TEMPLATE if not remove_qa_template else PROMPT_TEMPLATE + return prompt.format(context=context, question=q) def augment_prompt(sample, closed_qa_prob, cot_prob): @@ -206,7 +209,11 @@ def get_dataset_configs( def create_messages( - ctxs: list[str], questions: list[list[str]], vllm_model: str, system_template: str + ctxs: list[str], + questions: list[list[str]], + vllm_model: str, + system_template: str, + remove_qa_template: bool, ) -> list[list[dict]]: """Create chat messages for the model.""" # if "gemma" in vllm_model: @@ -215,7 +222,9 @@ def create_messages( [ { "role": "user", - "content": (system_template + "\n\n\n" + get_prompt(ctx, q)).strip(), + "content": ( + system_template + "\n\n\n" + get_prompt(ctx, q, remove_qa_template) + ).strip(), } ] for ctx, q_list in zip(ctxs, questions) @@ -326,6 +335,8 @@ def self_generate( tk = get_tokenizer(args.vllm_model, train=True) self_qa_intx_tokens = tk(SELF_QA_INTX, add_special_tokens=False)["input_ids"] + if args.remove_qa_template: + self_qa_intx_tokens = tk("\n\n", 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) @@ -390,7 +401,11 @@ def self_generate( 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 + chunk_ctxs, + chunk_questions, + args.vllm_model, + SELF_GEN_SYSTEM_MSG, + args.remove_qa_template, ) if do_truncate: @@ -404,7 +419,7 @@ def self_generate( truncate_middle_if_too_long( ids, max_length=MODEL_CTX_LEN[args.vllm_model], - max_new_tokens=1024, + max_new_tokens=args.max_new_tokens, ) for ids in tokenized_contents["input_ids"] ] @@ -455,7 +470,7 @@ def execute_qa_generation( completions = llm.chat( messages, sampling_params=SamplingParams( - max_tokens=1024, + max_tokens=args.max_new_tokens, logprobs=k, temperature=temp, seed=42, @@ -697,6 +712,17 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Truncate contexts to fit model context length", ) + parser.add_argument( + "--remove_qa_template", + action="store_true", + help="Remove QA template formatting from prompts", + ) + parser.add_argument( + "--max_new_tokens", + type=int, + default=256, + help="Maximum number of new tokens to generate (default: 256)", + ) return parser.parse_args() @@ -737,7 +763,7 @@ if __name__ == "__main__": for ds_name, split in dataset_configs: print(f"Processing dataset: {ds_name}, split: {split}") self_generate( - ds_name, split, args, llm, SELF_GEN_SYSTEM_MSG, args.do_truncate + ds_name, split, args, llm, SELF_GEN_SYSTEM_MSG, None, args.do_truncate ) else: assert args.glob_pattern, ( diff --git a/scripts/toy_exp.sh b/scripts/toy_exp.sh index 6f1b7cb..5cbcb25 100644 --- a/scripts/toy_exp.sh +++ b/scripts/toy_exp.sh @@ -5,7 +5,13 @@ uv run data/generate_ctx_magic_num.py # WANDB_PROJECT=ctx-magic-num srun --partition=aiscilow --gpus=1 --unbuffered uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=32 --per_device_eval_batch_size=16 --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=100 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=2048 --max_packed_ctx_len=2048 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=-1 --max_val_samples_per_ds=100 --seed=1 # WANDB_PROJECT=ctx-magic-num srun --partition=aiscilow --gpus=1 --unbuffered uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=32 --per_device_eval_batch_size=16 --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=100 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=2048 --max_packed_ctx_len=2048 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=512 --min_ctx_chunk_len=25 --num_chunk_probs='{"1":"0.5", "2":"0.5"}' --max_val_samples_per_ds=100 --seed=1 # WANDB_PROJECT=ctx-magic-num srun --partition=aiscilow --gpus=1 --unbuffered uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=32 --per_device_eval_batch_size=16 --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=100 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=2048 --max_packed_ctx_len=2048 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=512 --min_ctx_chunk_len=25 --num_chunk_probs='{"1":"0.5", "2":"0.25", "3":"0.125", "4":"0.125"}' --max_val_samples_per_ds=100 --seed=1 -WANDB_PROJECT=ctx-magic-num run uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=16 --per_device_eval_batch_size=16 --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=100 --logging_steps=10 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=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 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=512 --min_ctx_chunk_len=25 --num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' --max_val_samples_per_ds=100 --seed=1 +WANDB_PROJECT=ctx-magic-num run uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=16 --per_device_eval_batch_size=16 --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=100 --logging_steps=10 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=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 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=512 --min_ctx_chunk_len=25 --num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' --max_val_samples_per_ds=100 --seed=1 --use_per_ctx_average_loss=True + +# self-distill version +# ce +# WANDB_PROJECT=ctx-magic-num run uv run train.py configs/toy_exp/ctx_magic_number_32_256_self_gen.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=16 --per_device_eval_batch_size=16 --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=100 --logging_steps=10 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=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=per_layer_activations --n_latent_queries=8 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=-1 --max_val_samples_per_ds=100 --seed=1 --notes="1 l1norm + ce loss + 1 chunk" --use_per_ctx_average_loss=False +# kl +# WANDB_PROJECT=ctx-magic-num run uv run train.py configs/toy_exp/ctx_magic_number_32_256_self_gen.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=16 --per_device_eval_batch_size=16 --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=100 --logging_steps=10 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=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=per_layer_activations --n_latent_queries=8 --use_kl_loss=True --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=-1 --max_val_samples_per_ds=100 --seed=1 --notes="1 l1norm + kl loss + 1 chunk" --use_per_ctx_average_loss=False # eval WANDB_MODE=disabled srun --partition=aiscilow --gpus=1 --unbuffered uv run run_eval.py --checkpoint_path CHECKPOINT_PATH --datasets ctx_magic_number_32_1024 ctx_magic_number_1024_2048 ctx_magic_number_2048_3072 ctx_magic_number_3072_4096 ctx_magic_number_4096_5120 ctx_magic_number_5120_6144 ctx_magic_number_6144_7168 ctx_magic_number_7168_8192 ctx_magic_number_8192_9216 ctx_magic_number_9216_10240 ctx_magic_number_10240_11264 ctx_magic_number_11264_12288 ctx_magic_number_12288_13312 ctx_magic_number_13312_14336 ctx_magic_number_14336_15360 ctx_magic_number_15360_16384 ctx_magic_number_16384_20480 ctx_magic_number_20480_24576 ctx_magic_number_24576_28672 ctx_magic_number_28672_32768 ctx_magic_number_32768_40960 ctx_magic_number_40960_49152 ctx_magic_number_49152_57344 ctx_magic_number_57344_65536 ctx_magic_number_65536_73728 ctx_magic_number_73728_81920 ctx_magic_number_81920_90112 ctx_magic_number_90112_98304 ctx_magic_number_98304_106496 ctx_magic_number_106496_114688 ctx_magic_number_114688_122880 ctx_magic_number_122880_131072 --max_ctx_chunk_len=1024 --split test --eval_batch_size_gen=4 diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index ed3b61b..933f78f 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -428,7 +428,7 @@ class CtxTrainingArguments: metadata={"help": "Whether to use KL loss."}, ) use_per_ctx_average_loss: bool = field( - default=True, + default=False, metadata={"help": "Whether to use per-context average loss."}, ) gen_lora_l1_reg_coef: float = field( diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index 3ba4149..22809c7 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -1,7 +1,6 @@ import json import logging import os -import random from collections.abc import Callable from glob import glob from hashlib import sha256 @@ -508,6 +507,12 @@ def get_sft_prompt_formatting_fn( return_assistant_tokens_mask=True, return_dict=True, ) + + if tokenizer.name_or_path == "google/gemma-2-2b-it": + # HACK: add bos at the beginning (see `Known Issues` in README.md) + for tok_ids in tokens["input_ids"]: + tok_ids = [tokenizer.bos_token_id] + tok_ids + labels = [] 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)] diff --git a/src/ctx_to_lora/trainer.py b/src/ctx_to_lora/trainer.py index 43d9f38..be6cf13 100644 --- a/src/ctx_to_lora/trainer.py +++ b/src/ctx_to_lora/trainer.py @@ -153,6 +153,9 @@ class DistillationTrainer(ModulatedModelTrainer): label_pos = torch.where(labels != -100) outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True) + if "logprobs_vals" not in inputs: + return (torch.tensor(0.0), outputs) if return_outputs else torch.tensor(0.0) + target_logp = inputs.pop("logprobs_vals").squeeze(0) indices = inputs.pop("logprobs_indices").squeeze(0) diff --git a/watcher.py b/watcher.py index 719709a..a3e02b3 100644 --- a/watcher.py +++ b/watcher.py @@ -1,6 +1,7 @@ import itertools import os import time +from argparse import Namespace from glob import glob import wandb @@ -72,6 +73,7 @@ if __name__ == "__main__": run_dir = file.split("/checkpoint")[0] run_name = run_dir.split("/")[-1] print(f"Evaluating {file}") + args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml"))) curstep = int(file.split("checkpoint-")[1].split("/")[0]) wandb_kwargs = { "project": os.getenv("WANDB_PROJECT"), @@ -85,26 +87,28 @@ if __name__ == "__main__": # TODO: have to change this for bigger models eval_batch_size = 8 eval_batch_size_gen = 8 - try: - # metrics = run_eval( - # checkpoint_path=file, - # eval_batch_size=eval_batch_size, - # split="validation", - # generative=False, - # ) - metrics = {} - except FileNotFoundError as e: - print(f"Error evaluating {file}: {e}. The checkpoint might be deleted.") - continue + metrics = {} + + # try: + # # metrics = run_eval( + # # checkpoint_path=file, + # # eval_batch_size=eval_batch_size, + # # split="validation", + # # generative=False, + # # ) + # except FileNotFoundError as e: + # print(f"Error evaluating {file}: {e}. The checkpoint might be deleted.") + # continue try: gen_metrics = run_eval( checkpoint_path=file, - eval_batch_size=eval_batch_size_gen, split="validation", + eval_batch_size=eval_batch_size_gen, + max_ctx_chunk_len=args.max_ctx_chunk_len, generative=True, ) except FileNotFoundError as e: - print(f"Error evaluating {file}: {e}. The checkpoint might be deleted.") + print(f"The checkpoint might be deleted. Error evaluating {file}: {e}.") gen_metrics = {} file = "" metrics.update(gen_metrics)