From ee9039295e75cf5886d052a413223debb3eb039b Mon Sep 17 00:00:00 2001 From: Rujikorn Charakorn Date: Fri, 29 Aug 2025 11:55:24 +0900 Subject: [PATCH] LoRA separate bias param + merger + toy ds (#15) * separate bias + head weight init * smaller bias init + sum with 1/sqrt(combined_r) scaling + better lora merge logic * ok init * chunked ctx magic num working * configs * no transpose + only sum aggregation * fix cli * cli * 32-256 data * example cli + smaller dataset * toy ctx magic num ds size + quantize + remove head bias * configs removed * require lora_bias * Merge remote-tracking branch 'origin/main' into new-arch --- README.md | 30 ++- configs/toy_exp/ctx_kv_64_1024_self_gen.yaml | 20 -- configs/toy_exp/ctx_kv_64_128_self_gen.yaml | 12 -- configs/toy_exp/ctx_kv_64_256_self_gen.yaml | 14 -- configs/toy_exp/ctx_magic_number_32_256.yaml | 14 ++ .../toy_exp/ctx_numbers_64_128_self_gen.yaml | 12 -- .../toy_exp/ctx_numbers_64_2048_self_gen.yaml | 28 --- data/generate_ctx_magic_number.py | 185 ++++++++++++----- run_eval.py | 11 +- src/ctx_to_lora/configs.py | 16 +- src/ctx_to_lora/data/definitions.py | 3 + src/ctx_to_lora/data/processing.py | 104 +++++----- src/ctx_to_lora/eval_utils.py | 9 +- src/ctx_to_lora/model_loading.py | 2 +- src/ctx_to_lora/modeling/hypernet.py | 185 +++++++++-------- src/ctx_to_lora/modeling/lora_layer.py | 8 +- src/ctx_to_lora/modeling/lora_merger.py | 188 ++++-------------- train.py | 14 +- 18 files changed, 400 insertions(+), 455 deletions(-) delete mode 100644 configs/toy_exp/ctx_kv_64_1024_self_gen.yaml delete mode 100644 configs/toy_exp/ctx_kv_64_128_self_gen.yaml delete mode 100644 configs/toy_exp/ctx_kv_64_256_self_gen.yaml create mode 100644 configs/toy_exp/ctx_magic_number_32_256.yaml delete mode 100644 configs/toy_exp/ctx_numbers_64_128_self_gen.yaml delete mode 100644 configs/toy_exp/ctx_numbers_64_2048_self_gen.yaml diff --git a/README.md b/README.md index fdebb0f..74ce064 100644 --- a/README.md +++ b/README.md @@ -43,19 +43,24 @@ WANDB_MODE=disabled run uv run train.py configs/context_numbers_10_self_gen.yaml ### Synthetic ctx numbers ```bash -WANDB_MODE=disabled run uv run train.py configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_128.yaml --model_name_or_path=google/gemma-3-1b-it --num_train_epochs=3 --per_device_train_batch_size=-1 --gradient_accumulation_steps=2 --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=100 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --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=per_layer_activations --n_latent_queries=8 --use_kl_loss=False --eval_on_start=True +WANDB_MODE=disabled run uv run train.py configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_128.yaml --model_name_or_path=google/gemma-3-1b-it --num_train_epochs=3 --per_device_train_batch_size=-1 --gradient_accumulation_steps=2 --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=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.0 --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=per_layer_activations --n_latent_queries=8 --use_kl_loss=False --eval_on_start=True +``` + +ctx magic num +```bash + 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.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 ``` ### Squad only ```bash -WANDB_MODE=disabled uv run train.py configs/squad.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=16000 --max_packed_ctx_len=32000 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=10 -``` --> +WANDB_PROJECT=ctx-squad-test run uv run train.py configs/squad.yaml --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 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --ctx_encoder_type=per_layer_activations --n_latent_queries=8 --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=1 --use_sequence_packing=True --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=10 --max_ctx_chunk_len=-1 +``` ### Synthetic data generation @@ -66,8 +71,8 @@ uv run data/generate_fav_num.py # new num repeat uv run data/generate_ctx_numbers.py -# kv (8-char uuid keys, 4-digit values) -# uv run data/generate_ctx_kv.py --key-type uuid --value-type digits +# ctx magic num (NIAH) +uv run data/generate_ctx_magic_number.py --tokenizer-name google/gemma-2-2b-it ``` Self-gen for the number toy dataset @@ -213,19 +218,28 @@ run python train.py configs/...yaml ... --from_pretrained_checkpoint=train_outpu ``` ### Evaluation +ctx magic num +```bash +# base model +WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --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 --split test --eval_batch_size_gen=16 + +# hypernet w/ 1024 max chunk size +WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/Aug26_05-46-31_slurm0-aiscinodeset-1_81810_f78c9c91/checkpoint-1000/pytorch_model.bin --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 --max_ctx_chunk_len=1024 --split test +``` + LongBench ```bash # generative WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/.../pytorch_model.bin --datasets negative_nq triviaqa_retrieved squad longbench_e --split test # hypernet checkpoint -WANDB_MODE=disabled run uv run run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad ropes drop longbench/gov_report_e longbench/multifieldqa_en_e longbench/2wikimqa_e --split test --max_ctx_chunk_len -1 --lora_aggregation sum --eval_batch_size_gen 8 +WANDB_MODE=disabled run uv run run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad ropes drop longbench/gov_report_e longbench/multifieldqa_en_e longbench/2wikimqa_e --split test --max_ctx_chunk_len -1 --eval_batch_size_gen 8 # squad only WANDB_MODE=disabled run uv run run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad --split test # chunking -WANDB_MODE=disabled run uv run run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad --split validation --max_ctx_chunk_len 100 --max_val_samples_per_ds 10 --lora_aggregation mean +WANDB_MODE=disabled run uv run run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad --split validation --max_ctx_chunk_len 100 --max_val_samples_per_ds 10 # base model WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --eval_batch_size 2 diff --git a/configs/toy_exp/ctx_kv_64_1024_self_gen.yaml b/configs/toy_exp/ctx_kv_64_1024_self_gen.yaml deleted file mode 100644 index c6fd62e..0000000 --- a/configs/toy_exp/ctx_kv_64_1024_self_gen.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# 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_0.0/ctx_kv_64_128 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_kv_128_256 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_kv_256_512 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_kv_512_768 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_kv_768_1024 - -val_ds_names: -- ctx_kv_64_128 -- ctx_kv_128_256 -- ctx_kv_256_512 -- ctx_kv_512_768 -- ctx_kv_768_1024 diff --git a/configs/toy_exp/ctx_kv_64_128_self_gen.yaml b/configs/toy_exp/ctx_kv_64_128_self_gen.yaml deleted file mode 100644 index 4394796..0000000 --- a/configs/toy_exp/ctx_kv_64_128_self_gen.yaml +++ /dev/null @@ -1,12 +0,0 @@ -# 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_0.0/ctx_kv_64_128 - -val_ds_names: -- ctx_kv_64_128 diff --git a/configs/toy_exp/ctx_kv_64_256_self_gen.yaml b/configs/toy_exp/ctx_kv_64_256_self_gen.yaml deleted file mode 100644 index 1341e8e..0000000 --- a/configs/toy_exp/ctx_kv_64_256_self_gen.yaml +++ /dev/null @@ -1,14 +0,0 @@ -# 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_0.0/ctx_kv_64_128 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_kv_128_256 - -val_ds_names: -- ctx_kv_64_128 -- ctx_kv_128_256 diff --git a/configs/toy_exp/ctx_magic_number_32_256.yaml b/configs/toy_exp/ctx_magic_number_32_256.yaml new file mode 100644 index 0000000..722f2c7 --- /dev/null +++ b/configs/toy_exp/ctx_magic_number_32_256.yaml @@ -0,0 +1,14 @@ +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: +- ctx_magic_number_32_128 +- ctx_magic_number_128_256 + +val_ds_names: +- ctx_magic_number_32_128 +- ctx_magic_number_128_256 diff --git a/configs/toy_exp/ctx_numbers_64_128_self_gen.yaml b/configs/toy_exp/ctx_numbers_64_128_self_gen.yaml deleted file mode 100644 index 47c88e0..0000000 --- a/configs/toy_exp/ctx_numbers_64_128_self_gen.yaml +++ /dev/null @@ -1,12 +0,0 @@ -# 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_0.0/ctx_numbers_64_128 - -val_ds_names: -- ctx_numbers_64_128 diff --git a/configs/toy_exp/ctx_numbers_64_2048_self_gen.yaml b/configs/toy_exp/ctx_numbers_64_2048_self_gen.yaml deleted file mode 100644 index f1d648d..0000000 --- a/configs/toy_exp/ctx_numbers_64_2048_self_gen.yaml +++ /dev/null @@ -1,28 +0,0 @@ -# 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_0.0/ctx_numbers_64_128 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_numbers_128_256 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_numbers_256_512 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_numbers_512_768 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_numbers_768_1024 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_numbers_1024_1280 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_numbers_1280_1536 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_numbers_1536_1792 -- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_numbers_1792_2048 - -val_ds_names: -- ctx_numbers_64_128 -- ctx_numbers_128_256 -- ctx_numbers_256_512 -- ctx_numbers_512_768 -- ctx_numbers_768_1024 -- ctx_numbers_1024_1280 -- ctx_numbers_1280_1536 -- ctx_numbers_1536_1792 -- ctx_numbers_1792_2048 \ No newline at end of file diff --git a/data/generate_ctx_magic_number.py b/data/generate_ctx_magic_number.py index 5510c00..f2a871b 100644 --- a/data/generate_ctx_magic_number.py +++ b/data/generate_ctx_magic_number.py @@ -8,7 +8,9 @@ import random # Config knobs (edit or use CLI) # ----------------------------- TOKENS_PER_BLOCK = 40 # rough heuristic tokens per noise block -BASE_SAMPLES_PER_BIN = 12_800 +BASE_SAMPLES_PER_BIN = ( + 320_000 # training samples budget scaler only (val/test fixed at 1000 each) +) RNG_SEED = 42 NOISE_BLOCK = "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again." SPECIAL_TPL = "The special magic number is {magic_number}." @@ -82,29 +84,19 @@ def _build_example(total_blocks: int, depth_bin: int) -> dict: return {"context": context, "prompt": prompt, "response": response} -def generate_magic_dataset(n: int, k: int) -> tuple[list[dict], list[dict], list[dict]]: - """Generate n samples for a given block length k, evenly distributed across 10 depth bins.""" - # Evenly divide n across 10 depth bins +def generate_examples(n: int, k: int) -> list[dict]: + """Generate n examples (all for block length k) evenly across 10 depth bins.""" + if n <= 0: + return [] base = n // 10 rem = n % 10 counts = [base + (1 if i < rem else 0) for i in range(10)] - - dataset: list[dict] = [] + out: list[dict] = [] for depth_bin, c in enumerate(counts): for _ in range(c): - dataset.append(_build_example(total_blocks=k, depth_bin=depth_bin)) - - # Randomize before splitting to ensure val/test are random samples - random.shuffle(dataset) - - # 98/1/1 split - total = len(dataset) - train_sz = int(0.98 * total) - val_sz = int(0.01 * total) - train = dataset[:train_sz] - val = dataset[train_sz : train_sz + val_sz] - test = dataset[train_sz + val_sz :] - return train, val, test + out.append(_build_example(total_blocks=k, depth_bin=depth_bin)) + random.shuffle(out) + return out def main(): @@ -112,11 +104,17 @@ def main(): description="Generate noise-wrapped special magic number dataset (similar structure to generate_ctx_kv.py)", ) parser.add_argument("--seed", type=int, default=RNG_SEED, help="Random seed") + parser.add_argument( + "--tokenizer-name", + type=str, + default="google/gemma-2-2b-it", + help=("Tokenizer name"), + ) parser.add_argument( "--base-samples-per-bin", type=int, default=BASE_SAMPLES_PER_BIN, - help="Baseline number of samples per token bin (scaled by bin width)", + help="Baseline number of TRAINING samples per token bin (scaled by bin width). Validation & test are always 1000 each.", ) parser.add_argument( "--out-prefix", @@ -148,47 +146,140 @@ def main(): random.seed(args.seed) - # Same token bins pattern as generate_ctx_kv.py (current version) - tok_bins = [(64, 128), (128, 256), (256, 512)] + [ - (512 + 256 * i, 512 + 256 * (i + 1)) for i in range(14) - ] + # ---------------------------------------------------- + # Optional: report tokenizer-based token length stats + # ---------------------------------------------------- + if args.tokenizer_name: + try: + from transformers import AutoTokenizer # type: ignore + except Exception as e: # pragma: no cover + raise RuntimeError( + "Failed to import transformers. Install it or omit --tokenizer-name." + ) from e + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=True) + noise_token_count = len(tokenizer(NOISE_BLOCK).input_ids) + special_example = SPECIAL_TPL.format(magic_number="0000") + special_token_count = len(tokenizer(special_example).input_ids) + print( + f"[Tokenizer: {args.tokenizer_name}] Noise block tokens: {noise_token_count} | Special line tokens: {special_token_count}" + ) + + tok_bins = [(32, 128), (128, 256), (256, 512), (512, 1024), (32, 1024)] + [ + (1024 * i, 1024 * (i + 1)) for i in range(1, 16) + ] if args.only_first_n_bins is not None: tok_bins = tok_bins[: args.only_first_n_bins] - # Map token bins to block-length bins using heuristic - len_bins = [ - (lo // args.tokens_per_block, hi // args.tokens_per_block) - for (lo, hi) in tok_bins - ] + if args.tokenizer_name: + max_hi = max(hi for _, hi in tok_bins) + + def measure_len(k: int) -> int: + if k == 1: + ctx = SPECIAL_TPL.format(magic_number="0000") + else: + blocks = [NOISE_BLOCK] * (k - 1) + [ + SPECIAL_TPL.format(magic_number="0000") + ] + ctx = SEP.join(blocks) + return len(tokenizer(ctx).input_ids) + + lengths: list[int] = [0] + k = 1 + while True: + L = measure_len(k) + lengths.append(L) + if L >= max_hi: + break + k += 1 + + len_bins = [] + for lo, hi in tok_bins: + k_lo = None + for kk in range(1, len(lengths)): + if lengths[kk] >= lo: + k_lo = kk + break + if k_lo is None or lengths[k_lo] >= hi: + len_bins.append((0, 0)) + continue + k_hi = len(lengths) + for kk in range(k_lo, len(lengths)): + if lengths[kk] >= hi: + k_hi = kk + break + len_bins.append((k_lo, k_hi)) + + base_tokens = lengths[1] + delta = (lengths[2] - lengths[1]) if len(lengths) > 2 else 0 + print( + f"Using tokenizer-measured block ranges. base_tokens={base_tokens} approx_delta={delta}" + ) + else: + len_bins = [ + (lo // args.tokens_per_block, hi // args.tokens_per_block) + for (lo, hi) in tok_bins + ] if args.dry_run: - k = max(1, len_bins[0][0] or 1) - train, _, _ = generate_magic_dataset(n=10, k=k) - print("Sample entry:") - print(json.dumps(train[0], indent=2)) + for lb in len_bins: + if lb[1] > lb[0]: + k = max(1, lb[0]) + sample = generate_examples(10, k) + print("Sample entry:") + print(json.dumps(sample[0], indent=2)) + break return - + # ----------------------------------------------- + # Main generation per token bin + # ----------------------------------------------- + TARGET_VAL = 1000 + TARGET_TEST = 1000 for len_bin, tok_bin in zip(len_bins, tok_bins): - bin_size = max(1, len_bin[1] - len_bin[0]) + if len_bin[1] <= len_bin[0]: + print(f"Skipping token bin {tok_bin} (no valid block counts)") + continue + k_start = max(1, len_bin[0]) + k_end = max(1, len_bin[1]) + k_values = list(range(k_start, k_end)) + bin_size = len(k_values) save_dir = f"{args.out_prefix}_{tok_bin[0]}_{tok_bin[1]}" - train_data: list[dict] = [] + training_enabled = tok_bin[1] <= 512 # unchanged policy + if training_enabled: + train_data: list[dict] = [] + # Distribute training budget across k values. + # Scale: per_k = base_samples_per_bin / bin_size + per_k_train = max(1, args.base_samples_per_bin // max(1, bin_size)) + for k in k_values: + train_data += generate_examples(per_k_train, k) val_data: list[dict] = [] test_data: list[dict] = [] - - for k in range(max(1, len_bin[0]), max(1, len_bin[1])): - per_k = max(1, args.base_samples_per_bin // bin_size) - train, val, test = generate_magic_dataset(n=per_k, k=k) - train_data += train - val_data += val - test_data += test - + base_val = TARGET_VAL // bin_size + rem_val = TARGET_VAL % bin_size + base_test = TARGET_TEST // bin_size + rem_test = TARGET_TEST % bin_size + for idx, k in enumerate(k_values): + n_val_k = base_val + (1 if idx < rem_val else 0) + n_test_k = base_test + (1 if idx < rem_test else 0) + if n_val_k: + val_data += generate_examples(n_val_k, k) + if n_test_k: + test_data += generate_examples(n_test_k, k) + random.shuffle(val_data) + random.shuffle(test_data) os.makedirs(save_dir, exist_ok=True) - save_jsonl(train_data, f"{save_dir}/train.jsonl") + if training_enabled: + save_jsonl(train_data, f"{save_dir}/train.jsonl") save_jsonl(val_data, f"{save_dir}/val.jsonl") save_jsonl(test_data, f"{save_dir}/test.jsonl") - - print(f"Dataset generated and saved at {save_dir}") + if training_enabled: + print( + f"Dataset generated at {save_dir} (train={len(train_data)} val={len(val_data)} test={len(test_data)})" + ) + else: + print( + f"Dataset (val/test only) generated at {save_dir} (val={len(val_data)} test={len(test_data)})" + ) if __name__ == "__main__": diff --git a/run_eval.py b/run_eval.py index 8c59218..58446a5 100644 --- a/run_eval.py +++ b/run_eval.py @@ -66,12 +66,6 @@ if __name__ == "__main__": default=-1, help="Maximum length of context chunk for evaluation", ) - parser.add_argument( - "--lora_aggregation", - choices=["mean", "sum"], - default="sum", - help="LoRA aggregation method", - ) parser.add_argument( "--max_new_tokens", type=int, @@ -87,6 +81,11 @@ if __name__ == "__main__": cli_args = vars(parser.parse_args()) # setup_logging(output_dir, debug=os.getenv("DEBUG", False)) + if cli_args["model_name_or_path"]: + assert cli_args["max_ctx_chunk_len"] <= 0, ( + f"Evaluating base model shouldn't be used with `max_ctx_chunk_len`" + ) + eval_batch_size_gen = cli_args.pop("eval_batch_size_gen") eval_batch_size = cli_args.pop("eval_batch_size") # run_eval( diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index 6513706..c806c56 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -380,12 +380,9 @@ class CtxTrainingArguments: "they will be split up into multiple samples." }, ) - lora_aggregation: Literal["mean", "sum"] = field( - default="sum", metadata={"help": "LoRA aggregation method."} - ) - random_chunking: bool = field( - default=False, - metadata={"help": "Whether to use random chunking during training."}, + num_chunk_probs: dict = field( + default=None, + metadata={"help": "Probability distribution over chunk nums."}, ) max_ctx_chunk_len: int = field( default=-1, @@ -400,8 +397,8 @@ class CtxTrainingArguments: "help": "Min context chunk length. Used only with random chunking training" }, ) - max_ctx_chunk_num: int = field( - default=-1, + max_ctx_chunk_num: int | None = field( + default=None, metadata={"help": "Max number of context chunks per sample."}, ) max_packed_inp_len: int | None = field( @@ -534,6 +531,9 @@ class CtxEncoderArguments: "Default to L//4 where L is the number of layers of the ctx model" }, ) + quantize_ctx_encoder: bool = field( + default=False, metadata={"help": "Wheter to quantize the ctx encoder."} + ) @dataclass diff --git a/src/ctx_to_lora/data/definitions.py b/src/ctx_to_lora/data/definitions.py index 79f69f4..b8eb625 100644 --- a/src/ctx_to_lora/data/definitions.py +++ b/src/ctx_to_lora/data/definitions.py @@ -695,6 +695,9 @@ DS_KWARGS = { tok_bins = [(64, 128), (128, 256), (256, 512)] + [ (512 + 256 * i, 512 + 256 * (i + 1)) for i in range(14) ] +tok_bins += [(32, 128), (128, 256), (256, 512), (512, 1024), (32, 1024)] + [ + (1024 * i, 1024 * (i + 1)) for i in range(1, 16) +] for toy_ds_name in ["ctx_numbers", "ctx_kv", "ctx_magic_number"]: for tok_bin in tok_bins: DS_KWARGS[f"{toy_ds_name}_{tok_bin[0]}_{tok_bin[1]}"] = dict( diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index bc37e3e..3ba4149 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -7,6 +7,7 @@ from glob import glob from hashlib import sha256 from math import ceil, isclose from os import path +from random import choices from typing import Any import datasets @@ -210,7 +211,7 @@ def get_tokenized_dataset( ctx_tokenizer: PreTrainedTokenizerBase, max_ctx_chunk_len: int, min_ctx_chunk_len: int, - random_chunking: bool, + num_chunk_probs: dict[int, float] | None, max_ctx_chunk_num: int, add_ctx_to_chat: bool, use_kl_loss: bool, @@ -236,7 +237,7 @@ def get_tokenized_dataset( add_ctx_to_chat=add_ctx_to_chat, max_ctx_chunk_len=max_ctx_chunk_len, min_ctx_chunk_len=min_ctx_chunk_len, - random_chunking=random_chunking, + num_chunk_probs=num_chunk_probs, max_ctx_chunk_num=max_ctx_chunk_num, need_ctx_ids=need_ctx_ids, split=split, @@ -308,7 +309,7 @@ def construct_and_tokenize_ctx_qa( need_ctx_ids, max_ctx_chunk_len, min_ctx_chunk_len, - random_chunking, + num_chunk_probs, max_ctx_chunk_num, ds, split, @@ -376,16 +377,16 @@ def construct_and_tokenize_ctx_qa( # with some big caveats, e.g., losing order info split_ctx_kwargs = { "max_chunk_len": max_ctx_chunk_len, - "min_chunk_len": min_ctx_chunk_len - if min_ctx_chunk_len > 0 - else max_ctx_chunk_len, - "random_chunking": random_chunking, + "min_chunk_len": min_ctx_chunk_len, + "num_chunk_probs": num_chunk_probs, "max_num_split": max_ctx_chunk_num, "model_name_or_path": tokenizer.name_or_path, "is_train": is_train, } logging.info(f"Chunking context with {split_ctx_kwargs=}") - tokenized_ds = tokenized_ds.map(split_too_long_ctx, fn_kwargs=split_ctx_kwargs) + tokenized_ds = tokenized_ds.map( + split_too_long_ctx, fn_kwargs=split_ctx_kwargs, num_proc=16 + ) logging.info( f"Avg. num chunks per ctx: {np.mean(list(map(len, tokenized_ds['ctx_ids'])))}" ) @@ -405,14 +406,14 @@ def construct_and_tokenize_ctx_qa( if "train" not in split: # squeeze since we always have one query per sample in eval tokenized_ds = tokenized_ds.map(squeeze_tokens, num_proc=num_proc) - tokenized_ds = tokenized_ds.map( - truncate_middle_if_too_long, - fn_kwargs={ - "max_length": base_model_max_len, - "columns": ["input_ids", "labels"], - "max_new_tokens": max_new_tokens, - }, - ) + # tokenized_ds = tokenized_ds.map( + # truncate_middle_if_too_long, + # fn_kwargs={ + # "max_length": base_model_max_len, + # "columns": ["input_ids", "labels"], + # "max_new_tokens": max_new_tokens, + # }, + # ) tokenized_ds = tokenized_ds.map( add_length_info, fn_kwargs={"columns": ["input_ids"]}, @@ -575,11 +576,11 @@ def convert_ctx_prompt_response_to_messages( def split_too_long_ctx( sample: dict[str, Any], model_name_or_path: str, + num_chunk_probs: dict[int, float] | None, max_chunk_len: int, min_chunk_len: int, - max_num_split: int, + max_num_split: int | None, is_train: bool, - random_chunking: bool, ) -> dict[str, Any]: """ Split context into smaller chunks if it exceeds the maximum length. @@ -589,47 +590,52 @@ def split_too_long_ctx( max_chunk_len: Maximum length for each context chunk max_num_split: Maximum number of splits allowed - # Training - random_chunking: Wheter to use stochastic chunking with - `min_chunk_len` to `max_chunk_len` sizes - min_chunk_len: Minimum length for each context chunk (used when `random_chunking=True`) - Returns: Dictionary with split context data """ chunk_len = max_chunk_len - if is_train: - # TODO: for training, we might wanna sort the context by num chunks - # since merging a batch of chunked loras need padded in the rank axis - # e.g., ctx1 has 5 chunks (rank-48), ctx2 has 10 chunks (rank-88) + ctx_ids = sample["ctx_ids"] + # Early exits + if chunk_len <= 0 and max_num_split is None: + return {"ctx_ids": [ctx_ids]} - # e.g., even if the ctx is not too long, we still split it randomly? - # say 0-4k is max len for one split, - # for some ctx shorter than 4k, we leave it as is - # for some we split to smaller chunks - if random_chunking: - chunk_len = random.randint(min_chunk_len, max_chunk_len) + n_chunks = None # will be sampled (train) or derived (eval) + if is_train and num_chunk_probs is not None: + # Adjust theoretical upper bound based on min_chunk_len if provided + if min_chunk_len: + max_num_split = ceil(len(ctx_ids) / min_chunk_len) + # New logic: sample number of chunks from num_chunk_probs (after filtering) + + # Keep only feasible chunk counts <= max_num_split and > 0 + if max_num_split is not None: + filt = {k: v for k, v in num_chunk_probs.items() if 0 < k <= max_num_split} + else: + filt = num_chunk_probs + + # Ensure each chunk will not exceed max_chunk_len; enforce minimum required chunks + min_required = ( + max(1, ceil(len(ctx_ids) / max_chunk_len)) if max_chunk_len > 0 else 1 + ) + # Remove options that would yield chunk length > max_chunk_len + filt = {k: v for k, v in filt.items() if k >= min_required} + n_chunks = choices(list(filt.keys()), weights=list(filt.values()), k=1)[0] + + # Derive n_chunks if not sampled (e.g., eval or fallback) + if n_chunks is None: + n_chunks = ceil(len(ctx_ids) / chunk_len) + # Safety: at least 1 + n_chunks = max(1, n_chunks) + if n_chunks == 1: + return {"ctx_ids": [ctx_ids]} + + avg_len = ceil(len(ctx_ids) / n_chunks) + chunks = [ctx_ids[i : i + avg_len] for i in range(0, len(ctx_ids), avg_len)] ctx_affixes = CTX_AFFIXES[model_name_or_path] prefix = ctx_affixes["prefix"] suffix = ctx_affixes["suffix"] - ctx_ids = sample["ctx_ids"] - - if chunk_len <= 0 and max_num_split <= 0: - return {"ctx_ids": [ctx_ids]} - - if len(ctx_ids) <= chunk_len: - return {"ctx_ids": [ctx_ids]} - - # uniform chunking - n_chunks = ceil(len(ctx_ids) / chunk_len) - avg_len = ceil(len(ctx_ids) / n_chunks) - - # Split the context into smaller chunks - chunks = [ctx_ids[i : i + avg_len] for i in range(0, len(ctx_ids), avg_len)] - - # this would exceed the avg_len a bit + # Apply affixes chunks[0] = chunks[0] + suffix for i in range(1, len(chunks) - 1): chunks[i] = prefix + chunks[i] + suffix diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py index a3e3254..de82425 100644 --- a/src/ctx_to_lora/eval_utils.py +++ b/src/ctx_to_lora/eval_utils.py @@ -728,7 +728,6 @@ def evaluate( args: Namespace, split: str, max_ctx_chunk_len: int, - lora_aggregation: str, max_new_tokens: int, generative: bool, ) -> dict[str, dict]: @@ -789,8 +788,8 @@ def evaluate( max_qas_per_sample=1, max_ctx_chunk_len=max_ctx_chunk_len, min_ctx_chunk_len=-1, - random_chunking=False, - max_ctx_chunk_num=-1, + num_chunk_probs=None, + max_ctx_chunk_num=None, base_model_max_len=model.base_model.config.max_position_embeddings, tokenizer=tokenizer, ctx_model_max_len=ctx_model_max_len, @@ -868,8 +867,6 @@ def evaluate( if max_ctx_chunk_len > 0: model.generate = model.generate_with_multi_loras - print(f"Using {lora_aggregation} for aggregation generated LoRAs.") - model.generate = partial(model.generate, lora_aggregation=lora_aggregation) trainer_kwargs = { "model": model, @@ -922,7 +919,6 @@ def run_eval( eval_batch_size: int = 8, max_val_samples_per_ds: int = -1, max_ctx_chunk_len: int = -1, - lora_aggregation: str = "sum", remove_context: bool = False, max_new_tokens: int = 256, generative: bool = False, @@ -992,7 +988,6 @@ def run_eval( args, split, max_ctx_chunk_len, - lora_aggregation, max_new_tokens, generative=generative, ) diff --git a/src/ctx_to_lora/model_loading.py b/src/ctx_to_lora/model_loading.py index 672ad56..dad159b 100644 --- a/src/ctx_to_lora/model_loading.py +++ b/src/ctx_to_lora/model_loading.py @@ -174,7 +174,7 @@ def get_lora_config(model_dir, **kwargs): base_model_name_or_path=model_dir, task_type="CAUSAL_LM", lora_dropout=kwargs.get("lora_dropout", 0.0), - lora_alpha=r ** (3 / 2) * 2, + lora_alpha=2 / r**0.5, ) peft_conf_kwargs.update(kwargs) diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index 0aa9e1b..20d1cfb 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/src/ctx_to_lora/modeling/hypernet.py @@ -3,10 +3,11 @@ from collections.abc import Iterable from contextlib import contextmanager from dataclasses import dataclass from functools import partial -from typing import Any, Literal +from math import sqrt +from typing import Any import torch -from einops import rearrange, unpack +from einops import unpack from einops.layers.torch import EinMix as Mix from jaxtyping import Float, Integer from peft import ( @@ -273,24 +274,26 @@ class HyperLoRA(nn.Module): self.layers = nn.Sequential(*layers) - d_lora = max(self.d_in[m] + self.d_out[m] for m in self.target_modules) + self.d_lora = max(self.d_in[m] + self.d_out[m] for m in self.target_modules) - # self.bias_a = nn.ParameterDict( - # { - # m: nn.Parameter( - # torch.normal( - # 0, 1 / self.d_in[m], (self.n_layers, self.r, self.d_in[m]) - # ) - # ) - # for m in self.target_modules - # } - # ) - # self.bias_b = nn.ParameterDict( - # { - # m: nn.Parameter(torch.zeros((self.n_layers, self.r, self.d_out[m]))) - # for m in self.target_modules - # } - # ) + self.bias_a = nn.ParameterDict( + { + m: nn.Parameter( + torch.normal( + 0, + 0.1 / (self.d_in[m] * self.r) ** 0.5, + (self.n_layers, self.r, self.d_in[m]), + ) + ) + for m in self.target_modules + } + ) + self.bias_b = nn.ParameterDict( + { + m: nn.Parameter(torch.zeros((self.n_layers, self.r, self.d_out[m]))) + for m in self.target_modules + } + ) if self.config.use_light_weight_lora: # light-weight lora projection (per layer, per module) @@ -358,13 +361,13 @@ class HyperLoRA(nn.Module): ) param.bias.data[:] = 0 - d_lora = self.config.light_weight_latent_size * 2 + self.d_lora = self.config.light_weight_latent_size * 2 self.d_in = {k: self.config.light_weight_latent_size for k in self.d_in} self.d_out = { k: self.config.light_weight_latent_size for k in self.d_out } - logger.info(f"Using light-weight LoRA with d_lora = {d_lora // 2}") + logger.info(f"Using light-weight LoRA with d_lora = {self.d_lora // 2}") n_modules = len(self.target_modules) # have to do this otherwise doesnt work with adamw_torch_fused @@ -378,23 +381,23 @@ class HyperLoRA(nn.Module): self.head = Mix( "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", weight_shape="n_layers d_latent d_lora", - # bias_shape=None, # no bias - bias_shape="n_layers r d_lora", + bias_shape=None, # no bias + # bias_shape="n_layers r d_lora", n_layers=len(self.layer_indices), d_latent=self.config.latent_size, r=self.config.lora_config.r, - d_lora=d_lora, + d_lora=self.d_lora, ) else: self.head = Mix( "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", weight_shape="n_layers d_latent d_lora", - # bias_shape=None, # no bias - bias_shape="n_layers d_lora", + bias_shape=None, # no bias + # bias_shape="n_layers d_lora", n_layers=len(self.layer_indices), d_latent=self.config.latent_size, r=self.config.lora_config.r, - d_lora=d_lora, + d_lora=self.d_lora, ) # else: # if self.config.use_per_rank_bias: @@ -425,25 +428,25 @@ class HyperLoRA(nn.Module): self.head = Mix( "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", weight_shape="n_layers n_modules d_latent d_lora", - # bias_shape=None, # no bias - bias_shape="n_layers n_modules r d_lora", + bias_shape=None, # no bias + # bias_shape="n_layers n_modules r d_lora", n_layers=len(self.layer_indices), n_modules=n_modules, d_latent=self.config.latent_size, r=self.config.lora_config.r, - d_lora=d_lora, + d_lora=self.d_lora, ) else: self.head = Mix( "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", weight_shape="n_layers n_modules d_latent d_lora", - # bias_shape=None, # no bias - bias_shape="n_layers n_modules d_lora", + bias_shape=None, # no bias + # bias_shape="n_layers n_modules d_lora", n_layers=len(self.layer_indices), n_modules=n_modules, d_latent=self.config.latent_size, r=self.config.lora_config.r, - d_lora=d_lora, + d_lora=self.d_lora, ) # else: # self.head = Mix( @@ -546,16 +549,16 @@ class HyperLoRA(nn.Module): # hidden_size=self.config.base_hidden_size, # ) - # def get_head_bias(self): - # bias_dict = dict() - # for module in self.target_modules: - # bias_A = self.bias_a[module][None, ...] - # bias_B = self.bias_b[module][None, ...] + def get_head_bias(self): + bias_dict = dict() + for module in self.target_modules: + bias_A = self.bias_a[module] + bias_B = self.bias_b[module] - # # transpose B - # bias_B = rearrange(bias_B, "bs n_layers r d_out -> bs n_layers d_out r") - # bias_dict[module] = dict(A=bias_A, B=bias_B) - # return bias_dict + # transpose B + # bias_B = rearrange(bias_B, "bs n_layers r d_out -> bs n_layers d_out r") + bias_dict[module] = dict(A=bias_A, B=bias_B) + return bias_dict def _to_lora_dict( self, flat_loras: Float[Tensor, "bs n_layers n_modules r max_io_dim"] @@ -581,7 +584,7 @@ class HyperLoRA(nn.Module): "bs n_layers r *", ) # transpose B - B = rearrange(B, "bs n_layers r d_out -> bs n_layers d_out r") + # B = rearrange(B, "bs n_layers r d_out -> bs n_layers d_out r") if self.config.use_light_weight_lora: # A = einsum( # self.pre_lora_projection[module], @@ -593,6 +596,7 @@ class HyperLoRA(nn.Module): # B, # "d_out d_latent, bs n_layers d_latent r -> bs n_layers d_out r", # ) + # this is broken? A = self.pre_lora_projection[module](A) B = self.post_lora_projection[module](B) @@ -617,16 +621,19 @@ class HyperLoRA(nn.Module): features: Float[Tensor, "bs seq_len feature_dim"], attn_mask: Integer[Tensor, "bs seq_len"] | None = None, position_ids: Integer[Tensor, "bs seq_len"] | None = None, + n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None, ): - # [bs, n_layers x n_total_modules x r, feature_dim] + # [bs, n_layers, n_total_modules, r, feature_dim] with torch.autocast(device_type="cuda", dtype=torch.bfloat16): - # OMG!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! lora_emb, extra_emb = self.aggregator(features, attn_mask, position_ids) - # lora_emb, extra_emb = unpack( - # emb, - # [[self.num_modules], [self.num_extra_modules]], - # "bs n_layers * feature_dim", - # ) + + # TODO: add pos emb + # TODO: + # concat queries from the same chunks + # do self-attn + # unpack back to [bs, n_layers, n_total_modules, r, feature_dim] + # here bs = sum(n_chunks) + # we then # [bs, n_layers, n_modules, r, max_in_d_outim] flat_loras = None @@ -662,7 +669,6 @@ class ModulatedPretrainedModel(nn.Module): use_base_input_as_ctx: bool = False, # need non-packed inputs for generation use_sequence_packing: bool = True, - lora_aggregation: Literal["mean", "sum"] = "sum", ): assert not use_base_input_as_ctx super().__init__() @@ -672,8 +678,6 @@ class ModulatedPretrainedModel(nn.Module): self.ctx_encoder_args = ctx_encoder_args self.use_base_input_as_ctx = use_base_input_as_ctx self.use_sequence_packing = use_sequence_packing - self.lora_aggregation = lora_aggregation - self.lora_id = 1 self.active_adapters = [] self.register_module("base_model", base_model) @@ -749,6 +753,7 @@ class ModulatedPretrainedModel(nn.Module): train=self.base_model.training, requires_grad=False, use_flash_attn=base_model_attn_impl == "flash_attention_2", + use_q_lora=self.ctx_encoder_args.quantize_ctx_encoder, ) self.ctx_encoder = CTX_ENCODER_CLS[self.ctx_encoder_args.ctx_encoder_type]( encoder_model, self.ctx_encoder_args @@ -776,30 +781,53 @@ class ModulatedPretrainedModel(nn.Module): self.base_model, self.hypernet.lora_config ) logger.debug(f"peft_weights: {peft_weights}") - self.hypernet.head.weight.data[:] = 0 - self.hypernet.head.bias.data[:] = 0 + # self.hypernet.head.weight.data[:] = 0 + r = self.hypernet_config.lora_config.r + nn.init.normal_( + self.hypernet.head.weight, + mean=0, + std=2 / sqrt(self.hypernet.config.latent_size + self.hypernet.d_lora), + # the head outputs per rank lora --> divide by r to scale down grad + ) + # nn.init.orthogonal_(self.hypernet.head.weight, gain=1.0) + # self.hypernet.head.bias.data[:] = 0 - for i, m in enumerate(self.hypernet.target_modules): - A = peft_weights[m]["lora_A"].weight.clone() # [r, in_d] - B = peft_weights[m]["lora_B"].weight.clone() # [out_d, r] - if self.hypernet.config.use_light_weight_lora: - A = A[:, : self.hypernet.config.light_weight_latent_size] - B = B[: self.hypernet.config.light_weight_latent_size] - if self.hypernet.config.per_rank_gen: - if not self.hypernet.config.use_per_rank_bias: - A = A[0:1] - B = B[:, 0:1] - biases = [A, B.T] + # for i, m in enumerate(self.hypernet.target_modules): + # w_shape = self.hypernet.head.weight[ + # ..., i, :, : self.hypernet.d_in[m] + # ].shape + # self.hypernet.head.weight[..., i, :, : self.hypernet.d_in[m]] = ( + # torch.normal( + # mean=0, + # std=1 + # / sqrt( + # self.hypernet.config.latent_size + self.hypernet.d_in[m] + # ), + # size=w_shape, + # ) + # ) + # A = peft_weights[m]["lora_A"].weight.clone() # [r, in_d] + # B = peft_weights[m]["lora_B"].weight.clone() # [out_d, r] + # if self.hypernet.config.use_light_weight_lora: + # A = A[:, : self.hypernet.config.light_weight_latent_size] + # B = B[: self.hypernet.config.light_weight_latent_size] + # if self.hypernet.config.per_rank_gen: + # if not self.hypernet.config.use_per_rank_bias: + # A = A[0:1] + # B = B[:, 0:1] + # biases = [A, B.T] - # bias-hyperinit - # init weights to zeros and bias to the base weights - bias_cat = torch.cat(biases, dim=1) - # scale the biases by the output size - # e.g., rank=16 gives bigger (2x?) gradient magnitudes at initialization - # compared to rank=8 - r = self.hypernet_config.lora_config.r - bias_cat = bias_cat / r - self.hypernet.head.bias.data[..., i, :, : bias_cat.shape[1]] = bias_cat + # # bias-hyperinit + # # init weights to zeros and bias to the base weights + # bias_cat = torch.cat(biases, dim=1) + # # scale the biases by the output size + # # e.g., rank=16 gives bigger (2x?) gradient magnitudes at initialization + # # compared to rank=8 + # r = self.hypernet_config.lora_config.r + # bias_cat = bias_cat / r + # self.hypernet.head.bias.data[..., i, :, : bias_cat.shape[1]] = bias_cat + + # self.hypernet.head.bias.requires_grad = False def state_dict(self, *args, **kwargs): # we assume ctx_encoder and base model is frozen here @@ -935,8 +963,7 @@ class ModulatedPretrainedModel(nn.Module): generated_loras = combine_lora( generated_loras, n_ctx_chunks, - aggregation=self.lora_aggregation, - # lora_bias=self.hypernet.get_head_bias(), + lora_bias=self.hypernet.get_head_bias(), ) # input_ids in model_inputs_kwargs contains only @@ -998,7 +1025,6 @@ class ModulatedPretrainedModel(nn.Module): ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None, n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None, n_queries: Integer[Tensor, "n_ctx"] | None = None, - lora_aggregation: Literal["mean", "sum"] = "sum", *model_inputs_args: Any, **model_inputs_kwargs: dict[str, Any], ): @@ -1009,8 +1035,7 @@ class ModulatedPretrainedModel(nn.Module): generated_loras = combine_lora( generated_loras, n_ctx_chunks, - aggregation=lora_aggregation, - # lora_bias=self.hypernet.get_head_bias(), + lora_bias=self.hypernet.get_head_bias(), ) # apply lora hook to the base model diff --git a/src/ctx_to_lora/modeling/lora_layer.py b/src/ctx_to_lora/modeling/lora_layer.py index 78c0b7e..863585a 100644 --- a/src/ctx_to_lora/modeling/lora_layer.py +++ b/src/ctx_to_lora/modeling/lora_layer.py @@ -16,7 +16,7 @@ def lora_forward( n_qs: Integer[Tensor, "n_ctx"], tot_q: int, A: Float[Tensor, "n_ctx r d_in"], - B: Float[Tensor, "n_ctx d_out r"], + B: Float[Tensor, "n_ctx r d_out"], lora_dropout_p: float, scaling: float, self, @@ -32,7 +32,7 @@ def lora_forward( 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 = einsum(B, delta_x, "tot_q r d_out, tot_q s_len r -> tot_q s_len d_out") delta_x = delta_x * scaling return (base_out + delta_x).to(base_out.dtype) @@ -44,7 +44,7 @@ def lora_forward_packed( seq_lens: Integer[Tensor, "tot_q"], tot_len: int, A: Float[Tensor, "n_ctx r d_in"], - B: Float[Tensor, "n_ctx d_out r"], + B: Float[Tensor, "n_ctx r d_out"], lora_dropout_p: float, scaling: float, self, @@ -69,7 +69,7 @@ def lora_forward_packed( repeated_A, delta_x, "tot_len r d_in, bs tot_len d_in -> bs tot_len r" ) delta_x = einsum( - repeated_B, delta_x, "tot_len d_out r, bs tot_len r -> bs tot_len d_out" + repeated_B, delta_x, "tot_len r d_out, bs tot_len r -> bs tot_len d_out" ) delta_x = delta_x * scaling diff --git a/src/ctx_to_lora/modeling/lora_merger.py b/src/ctx_to_lora/modeling/lora_merger.py index e51b8f1..b19b98e 100644 --- a/src/ctx_to_lora/modeling/lora_merger.py +++ b/src/ctx_to_lora/modeling/lora_merger.py @@ -2,8 +2,6 @@ Utilities for merging / aggregating LoRA adapters coming from multiple chunks. """ -from typing import Literal - import torch from einops import rearrange from jaxtyping import Integer @@ -17,173 +15,57 @@ def compute_rank(n_lora, rank): def combine_lora( generated_loras: dict[str, dict[str, Tensor]], n_chunks: Integer[Tensor, "n_ctx"], - aggregation: Literal["mean", "sum"], - lora_bias: dict[str, dict[str, Tensor]] = None, + lora_bias: dict[str, dict[str, Tensor]], ) -> dict[str, dict[str, Tensor]]: - """Combine per-chunk LoRA adapter weights into grouped higher-rank adapters. - - Parameters - ---------- - generated_loras - Mapping: module_name -> {'A': Tensor, 'B': Tensor} - Each tensor shape: [total_chunks, n_layers, r, dim] for A, and - [total_chunks, n_layers, dim, r] for B (but we treat generically here). - n_chunks - 1D tensor giving number of chunks in each aggregation group. len(n_chunks)=G. - aggregation - 'mean' or 'sum'. Controls scaling and bias handling. - lora_bias - Optional mapping mirroring generated_loras providing bias LoRA tensors (same - per-chunk shape except missing the leading chunk dimension which is expanded). - - Returns - ------- - dict - module_name -> {'A': combined_A, 'B': combined_B} - combined_A shape: [G, n_layers, max_rank, dim] - combined_B shape: [G, n_layers, dim, max_rank] - """ - # TODO: optimize! - # assert aggregation in ["mean", "sum"] - if (n_chunks == 1).all(): - return generated_loras - - # Precompute common quantities total_chunks = int(n_chunks.sum()) # Assume all modules share same base rank r - # (second-to-last dim index for A, last dim for B) - # Use first module's A tensor to infer rank first_module = next(iter(generated_loras)) base_rank = generated_loras[first_module]["A"].shape[-2] - max_rank_needed = max(compute_rank(n, base_rank) for n in n_chunks) + max_rank_needed = compute_rank(n_chunks.max(), base_rank) - # Initialize output container combined_loras: dict[str, dict[str, Tensor]] = { module: {"A": None, "B": None} for module in generated_loras.keys() } + rank_dim = 2 + num_groups = len(n_chunks) + rank_per_group = n_chunks * base_rank - # Iterate over modules and LoRA matrices A / B. - # For A we concatenate along its rank dimension index=2; for B along index=3. for module_name, module_loras in generated_loras.items(): - for matrix_key, rank_dim in (("A", 2), ("B", 3)): - # Shape conventions (both stored with leading chunk axis): - # loras: [total_chunks, n_layers, r, dim] (A) - # or [total_chunks, n_layers, dim, r] (B) + for matrix_key in ("A", "B"): + bias_tensor = lora_bias[module_name][matrix_key] loras = module_loras[matrix_key] - # For 'mean' aggregation, scale each chunk's LoRA by sqrt(group_size) - if aggregation == "mean": - # Repeat for each chunk within a group: e.g., - # if n_chunks=[2,3] -> repeat factors [sqrt2, sqrt2, sqrt3, sqrt3, sqrt3] - sqrt_group_sizes = n_chunks**0.5 # [G] - sqrt_group_sizes_per_chunk = sqrt_group_sizes.repeat_interleave( - n_chunks, dim=0 - ) - loras = loras / sqrt_group_sizes_per_chunk.view(total_chunks, 1, 1, 1) + flat_loras = rearrange( + loras, "tot_chunks n_layers r dim -> 1 n_layers (tot_chunks r) dim" + ) + per_group_deltas = flat_loras.split(rank_per_group.tolist(), dim=rank_dim) - # Split the leading chunk dimension into groups per n_chunks specification - per_group_list = loras.split(n_chunks.tolist(), dim=0) + combined_shape = [num_groups, *per_group_deltas[0].shape[1:]] + combined_shape[rank_dim] = max_rank_needed - # Concat each group - if matrix_key == "A": - rearrange_pattern = "chunks n_layers r dim -> 1 n_layers (chunks r) dim" - else: - rearrange_pattern = "chunks n_layers dim r -> 1 n_layers dim (chunks r)" - per_group_deltas = [ - rearrange(group_tensor, rearrange_pattern) - for group_tensor in per_group_list - ] - - bias_tensor = None - if lora_bias is not None: - bias_tensor = lora_bias[module_name][matrix_key] - per_group_deltas = [ - torch.cat([group_tensor, bias_tensor], dim=rank_dim) - for group_tensor in per_group_deltas - ] - - # Delegate to helper for combining rank expansion across groups. - combined_loras[module_name][matrix_key] = _combine_single_matrix( - per_group_deltas=per_group_deltas, - bias=bias_tensor, - n_chunks=n_chunks, - base_rank=base_rank, - max_rank_needed=max_rank_needed, - aggregation=aggregation, - matrix_key=matrix_key, - rank_dim=rank_dim, + combined = torch.zeros( + *combined_shape, + device=per_group_deltas[0].device, + dtype=per_group_deltas[0].dtype, ) + for g, deltas in enumerate(per_group_deltas): + combined_rank = deltas.shape[rank_dim] + + # Build slice pattern, slice up to combined_rank. + slice_pattern = [g, slice(None), slice(None), slice(None)] + slice_pattern[rank_dim] = slice(combined_rank) + + combined[slice_pattern] = deltas + + if bias_tensor is not None: + bias_slice_pattern = [g, slice(None), slice(None), slice(None)] + bias_slice_pattern[rank_dim] = slice( + combined_rank, combined_rank + base_rank + ) + combined[bias_slice_pattern] = bias_tensor + + combined_loras[module_name][matrix_key] = combined + return combined_loras - - -def _combine_single_matrix( - *, - per_group_deltas: list[Tensor], - bias: Tensor | None, - n_chunks: Tensor, - base_rank: int, - max_rank_needed: int, - aggregation: str, - matrix_key: str, - rank_dim: int, -) -> Tensor: - """Assemble a single LoRA matrix (either 'A' or 'B') across groups. - - Parameters - ---------- - loras - Tensor of shape (total_chunks, n_layers, r, dim) if 'A' else - (total_chunks, n_layers, dim, r) if 'B'. - per_group_deltas - A list of concat'd tensors - bias_tensor - Expanded bias tensor of shape (G, n_layers, r, dim) for 'A' or - (G, n_layers, dim, r) for 'B', or None. - n_chunks - 1D tensor (G,) with number of chunks per group. - base_rank - Original (per-LoRA) rank r. - max_rank_needed - Maximum rank capacity allocated for any group (accounts for bias when summing). - aggregation - 'mean' or 'sum'. For 'sum' we optionally append a bias-related term. - matrix_key - 'A' or 'B'. Determines which axis is the rank axis and bias handling sign. - rank_dim - Dimension along which to concatenate within each group's tensors. - - Returns - ------- - Tensor - Combined LoRA matrix with shape (G, n_layers, max_rank, dim) for 'A' - or (G, n_layers, dim, max_rank) for 'B'. - """ - # assert matrix_key in ("A", "B"), f"matrix_key must be 'A' or 'B', got {matrix_key}" - - num_groups = len(n_chunks) - combined_shape = [num_groups, *per_group_deltas[0].shape[1:]] - # Adjust the rank dimension size placeholder. - if matrix_key == "A": - rank_dim_index = -2 # (G, layers, rank, dim) - else: - rank_dim_index = -1 # (G, layers, dim, rank) - combined_shape[rank_dim_index] = max_rank_needed - - combined = torch.zeros( - *combined_shape, - device=per_group_deltas[0].device, - dtype=per_group_deltas[0].dtype, - ) - - for g, deltas in enumerate(per_group_deltas): - combined_rank = deltas.shape[rank_dim] - - # Build slice pattern, slice up to combined_rank. - slice_pattern = [g, slice(None), slice(None), slice(None)] - slice_pattern[rank_dim] = slice(combined_rank) - - combined[slice_pattern] = deltas - - return combined diff --git a/train.py b/train.py index 5cd2cc2..c891b90 100755 --- a/train.py +++ b/train.py @@ -196,10 +196,7 @@ def main(): " as the context encoder" ) model = ModulatedPretrainedModel( - base_model, - hypernet_config, - ctx_encoder_args, - lora_aggregation=ctx_args.lora_aggregation, + base_model, hypernet_config, ctx_encoder_args ) else: @@ -254,6 +251,11 @@ def main(): # set default chunk size to max length of the ctx encoder ctx_args.max_ctx_chunk_len = ctx_model_max_len + if ctx_args.num_chunk_probs is not None: + ctx_args.num_chunk_probs = { + int(k): float(v) for k, v in ctx_args.num_chunk_probs.items() + } + _get_tokenized_dataset = partial( get_tokenized_dataset, max_qas_len=ctx_args.max_qas_len, @@ -265,7 +267,7 @@ def main(): add_ctx_to_chat=add_ctx_to_chat, max_ctx_chunk_len=ctx_args.max_ctx_chunk_len, min_ctx_chunk_len=ctx_args.min_ctx_chunk_len, - random_chunking=ctx_args.random_chunking, + num_chunk_probs=ctx_args.num_chunk_probs, max_ctx_chunk_num=ctx_args.max_ctx_chunk_num, use_kl_loss=ctx_args.use_kl_loss, ) @@ -391,7 +393,7 @@ if __name__ == "__main__": os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" os.environ["TOKENIZERS_PARALLELISM"] = "true" os.environ["WANDB_DIR"] = ".wandb/" - os.environ["WANDB_PROJECT"] = "ctx_to_lora" + os.environ["WANDB_PROJECT"] = os.getenv("WANDB_PROJECT") or "ctx_to_lora" os.environ["WANDB_WATCH"] = "" os.environ["WANDB_CONSOLE"] = "off" os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"