toy exp w/ self-gen + use double bos for eval (following vllm 0.8.4)

.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.

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
-```
+```
+
+## :warning: Known Issues
+For the ICLR version, we use two `<bos>` 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 `<bos>` tokens (only the `input_ids` not `ctx_ids`)
+
+NOTE: previously we were using `vllm==0.8.5` which does not have the double `<bos>` 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 `<bos>` bug, degraded `vllm` to `0.8.4`.

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
+
+

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, (

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

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(

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)]

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)
 

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)