apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-23 17:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

LoRA separate bias param + merger + toy ds (#15)

Browse source 
* 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
This commit is contained in:
Rujikorn Charakorn 2025-08-29 11:55:24 +09:00 • committed by GitHub
parent 9d82014dc2
commit ee9039295e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 400 additions and 455 deletions
Download patch file Download diff file Expand all files Collapse all files

30
README.md
View file

@ -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
```
<!-- ### FWQA-v2 Level-0 Tiny
```bash
WANDB_MODE=disabled run uv run train.py configs/fw_qa_v2_level_0_tiny.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=64 --gradient_accumulation_steps=8 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=1 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --max_packed_inp_len=24000 --max_packed_ctx_len=48000 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=50
```
``` -->
### 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

20
configs/toy_exp/ctx_kv_64_1024_self_gen.yaml
View file

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

12
configs/toy_exp/ctx_kv_64_128_self_gen.yaml
View file

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

14
configs/toy_exp/ctx_kv_64_256_self_gen.yaml
View file

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

14
configs/toy_exp/ctx_magic_number_32_256.yaml Normal file
View file

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

12
configs/toy_exp/ctx_numbers_64_128_self_gen.yaml
View file

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

28
configs/toy_exp/ctx_numbers_64_2048_self_gen.yaml
View file

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

185
data/generate_ctx_magic_number.py
View file

@ -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__":

11
run_eval.py
View file

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

16
src/ctx_to_lora/configs.py
View file

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

3
src/ctx_to_lora/data/definitions.py
View file

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

104
src/ctx_to_lora/data/processing.py
View file

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

9
src/ctx_to_lora/eval_utils.py
View file

@ -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,
)

2
src/ctx_to_lora/model_loading.py
View file

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

185
src/ctx_to_lora/modeling/hypernet.py
View file

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

8
src/ctx_to_lora/modeling/lora_layer.py
View file

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

188
src/ctx_to_lora/modeling/lora_merger.py
View file

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

14
train.py
View file

@ -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"