Hypernetworks that update LLMs to remember factual information https://arxiv.org/abs/2602.15902
Find a file
2025-06-17 11:34:16 +00:00
assets cover + readme 2025-05-29 15:20:09 +00:00
chat_templates/google fix chat template (include chat eos in labels) 2025-06-10 03:08:27 +00:00
configs data exp + packing 2025-06-17 11:34:16 +00:00
data data scripts + readme 2025-06-06 09:48:42 +00:00
icae_v2 rearrange + move to uv (#1) 2025-05-27 21:18:15 +09:00
scripts data exp + packing 2025-06-17 11:34:16 +00:00
src/ctx_to_lora data exp + packing 2025-06-17 11:34:16 +00:00
webui backward compat import 2025-05-30 08:11:47 +00:00
.gitignore more organized eval + viz_utils (llm-comparator) 2025-06-13 08:17:21 +00:00
.pre-commit-config.yaml rearrange + move to uv (#1) 2025-05-27 21:18:15 +09:00
install.sh update installation 2025-06-03 07:03:28 +00:00
intx_sft.py data exp + packing 2025-06-17 11:34:16 +00:00
pyproject.toml data exp + packing 2025-06-17 11:34:16 +00:00
README.md more organized eval + viz_utils (llm-comparator) 2025-06-13 08:17:21 +00:00
run_eval.py watcher + partial refactor eval + debug train w/ pwc and hotpot 2025-06-02 09:02:20 +00:00
run_lm_eval.py rearrange + move to uv (#1) 2025-05-27 21:18:15 +09:00
setup.py rearrange + move to uv (#1) 2025-05-27 21:18:15 +09:00
uv.lock data exp + packing 2025-06-17 11:34:16 +00:00
watcher.py watcher + partial refactor eval + debug train w/ pwc and hotpot 2025-06-02 09:02:20 +00:00

Ctx-to-LoRA



Project doc

🚀 API Usage [WIP]

from ctx_to_lora.modeling import ModulatedPretrainedModel
model = ModulatedPretrainedModel.from_state_dict(...)

ctx_info = "..."
query = "..."

ctx_ids = model.ctx_encoder.tokenize(ctx_info)
input_ids = model.tokenize(query)
outputs = model.generate(ctx_ids, input_ids)
print(model.decode(outputs))

🏋️ Training

🔢 HyperLoRA w/ context_numbers_10

WANDB_MODE=disabled uv run python intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=meta-llama/Llama-3.2-1B-Instruct --num_train_epochs=100 --per_device_train_batch_size=64 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --per_rank_gen=True --per_layer_processing=True --decoder_depth=2 --seed=1 --gen_lora_l1_reg_coef=0 --use_token_mixing=False --lora_r=8 --bf16=True --tf32=True --dataloader_num_workers=8 --dataloader_prefetch_factor=8

PwC + Hotpot training (for testing/debugging)

WANDB_MODE=disabled uv run intx_sft.py configs/pwc_hotpot_qa.yaml \
--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=32 \
--gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=hyper_lora \
--aggregator_type=perceiver \
--target_modules=down_proj \
--num_self_attends_per_block=8 --num_latent_factor=2 \
--lora_r=8 \
--eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 \
--neftune_noise_alpha=5 --use_light_weight_lora=False \
--load_best_model_at_end=True --metric_for_best_model=eval_pwc_loss --add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True --max_packed_inp_len=20000 --max_packed_ctx_len=40000 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \

HyperLoRA w/ self-gen 3 mini

# self-gen sft
# 3_mini config
uv run python data/self_generate_qa.py \
--vllm_model=google/gemma-2-2b-it --config=configs/self_gen_3_mini.yaml


WANDB_MODE=disabled uv run python intx_sft.py configs/self_gen_3_mini.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=10 --per_device_train_batch_size=32 --gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=2 --lora_r=8 --eval_steps=5000 --save_steps=5000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1

HyperLoRA w/ fw-qa pretrain only

WANDB_MODE=disabled run uv run python intx_sft.py configs/fw_qa_pretrain_only.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=4 --gradient_accumulation_steps=8 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=4 --num_latent_factor=1 --lora_r=8 --eval_steps=100 --save_steps=100 --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 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1

Favourite numbers data generation

uv run python data/generate_fav_num.py
# or
uv run python data/generate_fav_num_big.py

Data Generation

# [WIP] there are other datasets where we used openAI to gen QA pairs

# 0. download fineweb_edu to `data/raw_datasets/fineweb_edu
uv run data/download_fineweb_edu.py

# ===
# 1 and 2 can be run in parallel (depends on step 0)

# 1. generate QA data (3 QAs for each context)
# 000_0000[0-1] to 000_0000[8-9] for small split
vllm_model=mistralai/Mistral-Small-3.1-24B-Instruct-2503 uv run python generate_fw_edu_qa_vllm.py "00*_*" 3

# 2. Augmenting pre-training data w/ paraphrasing
# (000_00000 to 000_00009 for small split)
vllm_model=mistralai/Mistral-Small-3.1-24B-Instruct-2503 uv run python data/generate_fw_edu_augment_vllm.py 000_00000 

# ===
# 3 and 4 can be run in parallel (depends on step 0 and 1)

# 3. Pre-training data (no GPU needed) 
# augment each context by concat'ing QAs generated in step 1 at the end
uv run python data/generate_pretrain_from_fw_qa.py

# 4. self-generate SFT data
# each context has 3 chat instances as training data
# the responses are not the same as the ones in the pretraining though
# Example commands using gemma-2-2b-it
# self-gen data for fw_qa
# 000_0000[0-1] to 000_0000[8-9] for small split
uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_3/00*_*'
# self-gen data for other ds listed in qa_no_fw.yaml
uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_no_fw.yaml


Continue from a checkpoint

run python intx_sft.py configs/...yaml ... --from_pretrained_checkpoint=train_outputs/runs/May09_16-25-35_slurm0-a3nodeset-4_59459_ea85a571/checkpoint-10000/pytorch_model.bin --resume_from_checkpoint=train_outputs/runs/May09_16-25-35_slurm0-a3nodeset-4_59459_ea85a571/checkpoint-10000

Evaluation

LongBench

# generative
WANDB_MODE=disabled uv run python 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 uv run python run_eval.py --checkpoint_path train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/checkpoint-105000/pytorch_model.bin --datasets negative_nq triviaqa_retrieved squad longbench_e --split test

# base model
WANDB_MODE=disabled uv run python 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

# base model w/o context
WANDB_MODE=disabled uv run python run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --remove_context

# # benchmark
# cd LongBench/LongBench

# run python pred_ctx_to_lora.py  --checkpoint_path ../../train_outputs/runs/Mar16_12-38-01_slurm0-a3nodeset-12_54818_32426662/checkpoint-136782/pytorch_model.bin

# run python eval_ctx_to_lora.py --model_name Mar16_12-38-01_slurm0-a3nodeset-12_54818_32426662/checkpoint-136782 --checkpoint_path ../../train_outputs/runs/Mar16_12-38-01_slurm0-a3nodeset-12_54818_32426662/checkpoint-136782/pytorch_model.bin

LLM-comparator

# install nvm
curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.3/install.sh | bash
nvm install 16
nvm use 16

git clone https://github.com/PAIR-code/llm-comparator.git
cd llm-comparator
npm install
npm run build
# running llm-comparator webui
npm run serve

# in another terminal
# run http server for file fetching

# cd back to root folder first
# taken from https://stackoverflow.com/a/79135787
alias srv='echo -e "from sys import argv as a\nfrom http.server import HTTPServer as H, SimpleHTTPRequestHandler as HH, test as t\nclass C(HH):\n def end_headers (self):\n  self.send_header(a[2],a[3])\n  HH.end_headers(self)\nt(C,H,port=int(a[1]))" | /usr/bin/env python3 -- - 8001 "Access-Control-Allow-Origin" "*"'
srv

# copy-paste the relative path from the root
# e.g., http://localhost:8001/train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/eval-results-105000/comparator.json