Ctx-to-LoRA

--- [Project doc](https://docs.google.com/document/d/1RCQDzlVU7YGoTwR84gLQfhxTv0RFnW6srQlqfCC2bvQ/edit?usp=sharing) ## 🚀 API Usage [WIP] ```python 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 ```bash 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) ```bash 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 ```bash # 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 ```bash 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 ```bash uv run python data/generate_fav_num.py # or uv run python data/generate_fav_num_big.py ``` ### Data Generation ```bash # [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) 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 uv run python data/generate_fw_edu_augment_vllm.py # === # 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 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 ```bash 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 ```bash # generative 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 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 uv run python run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench_e --split test # base model w/o context run uv run python run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets 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 ```