diff --git a/README.md b/README.md index 81cecff..61ab0cc 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,16 @@ KL loss ```bash WANDB_MODE=disabled run uv run train.py configs/context_numbers_10_self_gen.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --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=early_exit --n_latent_queries=208 --use_kl_loss=True ``` -### FWQA-v2 Level-0 Tiny + + +### 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 +``` + + + -### Favourite numbers data generation +### Synthetic data generation ```bash -uv run python data/generate_fav_num.py -# or -uv run python data/generate_fav_num_big.py +# old number repeat +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 ``` Self-gen for the number toy dataset ```bash -run uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names context_numbers_2_10 --split train -run uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names context_numbers_2_10 --split validation +# old numbers +run uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names context_numbers_2_10 --split train +run uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names context_numbers_2_10 --split validation + +# new ctx numbers +run uv run data/self_generate_qa.py --vllm_model google/gemma-3-1b-it --ds_names ctx_numbers_64_128 --split train --remove_qa_template --max_new_tokens 150 ``` ### SQuAD @@ -153,23 +171,23 @@ uv run data/download_fineweb_edu.py 1. Recursively generate more data! (depends on step 0) ```bash # run from 000 to 0013 -run uv run python data/generate_fw_edu_qa_v2.py --shard_pattern "000_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048; -run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it; -run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_1*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it; -run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_2*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it +run uv run data/generate_fw_edu_qa_v2.py --shard_pattern "000_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048; +run uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it; +run uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_1*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it; +run uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_2*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it ``` 2. Self-generated response QA data (depends on step 0 and 1) ```bash # Example commands using gemma-2-2b-it # self-gen data for fw_qa_v2 -uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_v2/min_0_to_2000/013*_level_3*' --closed_qa_prob 1.0 # or 0.0 +uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_v2/min_0_to_2000/013*_level_3*' --closed_qa_prob 1.0 # or 0.0 # val split -uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_v2/min_0_to_2000/*_level_0_val.parquet' +uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_v2/min_0_to_2000/*_level_0_val.parquet' # self-gen data for other ds listed in qa_short_ctx_self_gen_no_fw_qa.yaml -uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_short_ctx_self_gen_no_fw_qa.yaml +uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_short_ctx_self_gen_no_fw_qa.yaml ``` @@ -198,22 +216,22 @@ run python train.py configs/...yaml ... --from_pretrained_checkpoint=train_outpu LongBench ```bash # 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 +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 python 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 --lora_aggregation sum --eval_batch_size_gen 8 # squad only -WANDB_MODE=disabled run uv run python 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 +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 python 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 --lora_aggregation mean # 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 +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 # 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 +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 --remove_context # # benchmark # cd LongBench/LongBench diff --git a/chat_templates/google/gemma-2-9b-it.jinja b/chat_templates/google/gemma-2-9b-it.jinja index 23129b0..4b683e5 100644 --- a/chat_templates/google/gemma-2-9b-it.jinja +++ b/chat_templates/google/gemma-2-9b-it.jinja @@ -14,7 +14,7 @@ {%- if message['role'] == 'user' and loop.first and system_message is defined %} {{ '' + role + '\n' + system_message + '\n\n' + message['content'] | trim + '\n' }} {%- elif message['role'] == 'assistant' %} - {{ '' + role + '\n' }}{% generation %}{{ message['content'] | trim + '\n' }}{% endgeneration %} + {{ '' + role + '\n' }}{% generation %}{{ message['content'] + '' }}{% endgeneration %}{{ '\n' }} {%- else %} {{ '' + role + '\n' + message['content'] | trim + '\n' }} {%- endif %} diff --git a/chat_templates/google/gemma-3-1b-it.jinja b/chat_templates/google/gemma-3-1b-it.jinja index 1713892..345ff87 100644 --- a/chat_templates/google/gemma-3-1b-it.jinja +++ b/chat_templates/google/gemma-3-1b-it.jinja @@ -22,7 +22,7 @@ {%- if (message['role'] == 'user') -%} {{ '' + role + '\n' + (first_user_prefix if loop.first else "") + message['content'] | trim + '\n' }} {%- elif (message['role'] == 'assistant') -%} - {{ '' + role + '\n' }}{% generation %}{{ message['content'] | trim + '\n' }}{% endgeneration %} + {{ '' + role + '\n' }}{% generation %}{{ message['content'] + '' }}{% endgeneration %}{{ '\n' }} {%- endif -%} {%- endfor -%} {%- if add_generation_prompt -%} diff --git a/chat_templates/google/gemma-3-4b-it.jinja b/chat_templates/google/gemma-3-4b-it.jinja index 1713892..345ff87 100644 --- a/chat_templates/google/gemma-3-4b-it.jinja +++ b/chat_templates/google/gemma-3-4b-it.jinja @@ -22,7 +22,7 @@ {%- if (message['role'] == 'user') -%} {{ '' + role + '\n' + (first_user_prefix if loop.first else "") + message['content'] | trim + '\n' }} {%- elif (message['role'] == 'assistant') -%} - {{ '' + role + '\n' }}{% generation %}{{ message['content'] | trim + '\n' }}{% endgeneration %} + {{ '' + role + '\n' }}{% generation %}{{ message['content'] + '' }}{% endgeneration %}{{ '\n' }} {%- endif -%} {%- endfor -%} {%- if add_generation_prompt -%} diff --git a/configs/gemma-3-1b-it/toy_exp/context_numbers_10.yaml b/configs/gemma-3-1b-it/toy_exp/context_numbers_10.yaml new file mode 100644 index 0000000..4d825f8 --- /dev/null +++ b/configs/gemma-3-1b-it/toy_exp/context_numbers_10.yaml @@ -0,0 +1,57 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: google/gemma-3-1b-it +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 64 +per_device_eval_batch_size: 128 +max_new_tokens: 64 +gen_per_device_eval_batch_size: 128 +max_val_samples_per_ds: 500 +# optim: schedule_free_adamw +learning_rate: 0.0001 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +# LoRA +lora_r: 16 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: +- data/raw_datasets/context_numbers_2 +- data/raw_datasets/context_numbers_3 +- data/raw_datasets/context_numbers_4 +- data/raw_datasets/context_numbers_5 +- data/raw_datasets/context_numbers_6 +- data/raw_datasets/context_numbers_7 +- data/raw_datasets/context_numbers_8 +- data/raw_datasets/context_numbers_9 +- data/raw_datasets/context_numbers_10 + +val_ds_names: +- data/raw_datasets/context_numbers_2 +- data/raw_datasets/context_numbers_3 +- data/raw_datasets/context_numbers_4 +- data/raw_datasets/context_numbers_5 +- data/raw_datasets/context_numbers_6 +- data/raw_datasets/context_numbers_7 +- data/raw_datasets/context_numbers_8 +- data/raw_datasets/context_numbers_9 +- data/raw_datasets/context_numbers_10 diff --git a/configs/gemma-3-1b-it/toy_exp/ctx_kv_64_128.yaml b/configs/gemma-3-1b-it/toy_exp/ctx_kv_64_128.yaml new file mode 100644 index 0000000..02f024c --- /dev/null +++ b/configs/gemma-3-1b-it/toy_exp/ctx_kv_64_128.yaml @@ -0,0 +1,12 @@ +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: +- ctx_kv_64_128 + +val_ds_names: +- ctx_kv_64_128 diff --git a/configs/gemma-3-1b-it/toy_exp/ctx_kv_64_128_self_gen.yaml b/configs/gemma-3-1b-it/toy_exp/ctx_kv_64_128_self_gen.yaml new file mode 100644 index 0000000..bfe2203 --- /dev/null +++ b/configs/gemma-3-1b-it/toy_exp/ctx_kv_64_128_self_gen.yaml @@ -0,0 +1,12 @@ +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: +- self_gen/google/gemma-3-1b-it_temp_0.0_closed_qa_prob_0.0/ctx_kv_64_128 + +val_ds_names: +- ctx_kv_64_128 diff --git a/configs/gemma-3-1b-it/toy_exp/ctx_kv_64_256.yaml b/configs/gemma-3-1b-it/toy_exp/ctx_kv_64_256.yaml new file mode 100644 index 0000000..4f281e0 --- /dev/null +++ b/configs/gemma-3-1b-it/toy_exp/ctx_kv_64_256.yaml @@ -0,0 +1,14 @@ +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: +- ctx_kv_64_128 +- ctx_kv_128_256 + +val_ds_names: +- ctx_kv_64_128 +- ctx_kv_128_256 diff --git a/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_1024.yaml b/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_1024.yaml new file mode 100644 index 0000000..4eae29c --- /dev/null +++ b/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_1024.yaml @@ -0,0 +1,21 @@ +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: +- ctx_numbers_64_128 +- ctx_numbers_128_256 +- ctx_numbers_256_512 +- ctx_numbers_512_768 +- ctx_numbers_768_1024 + + +val_ds_names: +- ctx_numbers_64_128 +- ctx_numbers_128_256 +- ctx_numbers_256_512 +- ctx_numbers_512_768 +- ctx_numbers_768_1024 diff --git a/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_128.yaml b/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_128.yaml new file mode 100644 index 0000000..8dbc6b9 --- /dev/null +++ b/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_128.yaml @@ -0,0 +1,12 @@ +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: +- ctx_numbers_64_128 + +val_ds_names: +- ctx_numbers_64_128 diff --git a/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_128_self_gen.yaml b/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_128_self_gen.yaml new file mode 100644 index 0000000..05b8530 --- /dev/null +++ b/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_128_self_gen.yaml @@ -0,0 +1,12 @@ +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: +- self_gen/google/gemma-3-1b-it_temp_0.0_closed_qa_prob_0.0/ctx_numbers_64_128 + +val_ds_names: +- ctx_numbers_64_128 diff --git a/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_2048.yaml b/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_2048.yaml new file mode 100644 index 0000000..c9a1108 --- /dev/null +++ b/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_2048.yaml @@ -0,0 +1,29 @@ +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_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 + + +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 diff --git a/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_256.yaml b/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_256.yaml new file mode 100644 index 0000000..d002fec --- /dev/null +++ b/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_256.yaml @@ -0,0 +1,14 @@ +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: +- ctx_numbers_64_128 +- ctx_numbers_128_256 + +val_ds_names: +- ctx_numbers_64_128 +- ctx_numbers_128_256 diff --git a/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_512.yaml b/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_512.yaml new file mode 100644 index 0000000..6c5b5e6 --- /dev/null +++ b/configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_512.yaml @@ -0,0 +1,17 @@ +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: +- ctx_numbers_64_128 +- ctx_numbers_128_256 +- ctx_numbers_256_512 + + +val_ds_names: +- ctx_numbers_64_128 +- ctx_numbers_128_256 +- ctx_numbers_256_512 diff --git a/configs/toy_exp/ctx_kv_64_1024_self_gen.yaml b/configs/toy_exp/ctx_kv_64_1024_self_gen.yaml new file mode 100644 index 0000000..c6fd62e --- /dev/null +++ b/configs/toy_exp/ctx_kv_64_1024_self_gen.yaml @@ -0,0 +1,20 @@ +# 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 new file mode 100644 index 0000000..4394796 --- /dev/null +++ b/configs/toy_exp/ctx_kv_64_128_self_gen.yaml @@ -0,0 +1,12 @@ +# 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 new file mode 100644 index 0000000..1341e8e --- /dev/null +++ b/configs/toy_exp/ctx_kv_64_256_self_gen.yaml @@ -0,0 +1,14 @@ +# 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_numbers_64_128_self_gen.yaml b/configs/toy_exp/ctx_numbers_64_128_self_gen.yaml new file mode 100644 index 0000000..47c88e0 --- /dev/null +++ b/configs/toy_exp/ctx_numbers_64_128_self_gen.yaml @@ -0,0 +1,12 @@ +# 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 new file mode 100644 index 0000000..f1d648d --- /dev/null +++ b/configs/toy_exp/ctx_numbers_64_2048_self_gen.yaml @@ -0,0 +1,28 @@ +# 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_kv.py b/data/generate_ctx_kv.py new file mode 100644 index 0000000..9a6ed66 --- /dev/null +++ b/data/generate_ctx_kv.py @@ -0,0 +1,243 @@ +import argparse +import json +import os +import random +import uuid +from collections.abc import Iterable + +import wonderwords + +# ----------------------------- +# Simple, editable config knobs +# ----------------------------- +# You can change these defaults or pass CLI flags: --key-type / --value-type +DEFAULT_KEY_TYPE = "uuid" # choices: "uuid" | "digits" | "words" +DEFAULT_VALUE_TYPE = "uuid" # choices: "uuid" | "digits" | "words" +PAIR_SEP = "\n" # between pairs +KV_SEP = " : " # between key and value +TOKENS_PER_PAIR = 10 # rough heuristic used for bucketing by context length +BASE_SAMPLES_PER_BIN = 128_000 +RNG_SEED = 42 + +nouns = wonderwords.random_word._get_words_from_text_file("nounlist.txt") +adjs = wonderwords.random_word._get_words_from_text_file("adjectivelist.txt") +# verbs = wonderwords.random_word._get_words_from_text_file("verblist.txt") +words = [f"{adj}-{noun}" for adj in adjs for noun in nouns] +words = sorted(list(set(words))) + + +def save_jsonl(data: list[dict], filepath: str) -> None: + """Save data to a JSONL file.""" + parent_dir = os.path.dirname(filepath) + if parent_dir: + os.makedirs(parent_dir, exist_ok=True) + with open(filepath, "w") as f: + for entry in data: + json.dump(entry, f) + f.write("\n") + + +def _gen_uuid8() -> str: + return uuid.uuid4().hex[:8] + + +def _gen_digits6() -> str: + return f"{random.randint(0, 999_999):06d}" + + +def _make_generator(kind: str): + kind = kind.lower() + if kind == "uuid": + return _gen_uuid8 + if kind == "digits": + return _gen_digits6 + if kind == "words": + return _gen_word + raise ValueError(f"Unknown kind '{kind}', expected 'uuid', 'digits', or 'words'") + + +def _gen_digits4() -> str: + return f"{random.randint(0, 9_999):04d}" + + +def _gen_word() -> str: + return random.choice(words) + + +def _make_value_generator(kind: str): + """Like _make_generator, but when kind is 'digits' use 4 digits for values.""" + kind = kind.lower() + if kind == "uuid": + return _gen_uuid8 + if kind == "digits": + return _gen_digits4 + if kind == "words": + return _gen_word + raise ValueError(f"Unknown kind '{kind}', expected 'uuid', 'digits', or 'words'") + + +def _unique(seq: Iterable[str]) -> list[str]: + s = set() + out = [] + for x in seq: + if x not in s: + s.add(x) + out.append(x) + return out + + +def render_context(pairs: list[tuple[str, str]]) -> str: + """Render key-value pairs as a single context string. + + Example: "k1:v1 | k2:v2 | k3:v3" + """ + return PAIR_SEP.join([f"{k}{KV_SEP}{v}" for k, v in pairs]) + + +def generate_kv_dataset(n: int, k: int, key_type: str, value_type: str): + """Generate n examples. Each example has k key-value pairs in the context. + + - context: "k1:v1 | k2:v2 | ..." + - prompt: "What is the value of `{k}`?" + - response: "v" (exactly the value corresponding to the chosen k) + """ + key_gen = _make_generator(key_type) + val_gen = _make_value_generator(value_type) + + dataset = [] + + for _ in range(n): + # Generate unique keys to avoid ambiguity + keys: list[str] = [] + while len(keys) < k: + keys = _unique([key_gen() for _ in range(k * 2)]) # oversample, then dedupe + keys = keys[:k] + values = [val_gen() for _ in range(k)] + + pairs = list(zip(keys, values)) + + # Choose a random key from the context to ask about + q_key = random.choice(keys) + q_val = values[keys.index(q_key)] + + entry = { + "context": f"key{KV_SEP}value\n" + render_context(pairs), + "prompt": f'What is the value of key "{q_key}"? Reply with only the value.', + "response": q_val, + } + dataset.append(entry) + + # Split into train/val/test matching the same structure as generate_ctx_numbers.py + total_size = len(dataset) + train_size = int(0.98 * total_size) + val_size = int(0.01 * total_size) + + train_data = dataset[:train_size] + val_data = dataset[train_size : train_size + val_size] + test_data = dataset[train_size + val_size :] + return train_data, val_data, test_data + + +def main(): + parser = argparse.ArgumentParser( + description="Generate key-value context QA data with the same structure as generate_ctx_numbers.py", + ) + parser.add_argument( + "--key-type", + choices=["uuid", "digits", "words"], + default=DEFAULT_KEY_TYPE, + help="Type of keys to generate: 8-char uuid, 6-digit numbers, or words", + ) + parser.add_argument( + "--value-type", + choices=["uuid", "digits", "words"], + default=DEFAULT_VALUE_TYPE, + help="Type of values to generate: 8-char uuid, 4-digit numbers, or words", + ) + parser.add_argument("--seed", type=int, default=RNG_SEED, help="Random seed") + 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)", + ) + parser.add_argument( + "--out-prefix", + type=str, + default="data/raw_datasets/ctx_kv", + help="Output directory prefix (bin range will be appended)", + ) + parser.add_argument( + "--tokens-per-pair", + type=int, + default=TOKENS_PER_PAIR, + help="Heuristic tokens per key:value pair for bucketing", + ) + parser.add_argument( + "--only-first-n-bins", + type=int, + default=None, + help="For quick tests: only generate the first N token bins", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print a small sample and exit without writing files", + ) + + args = parser.parse_args() + + random.seed(args.seed) + + # Same token bins as generate_ctx_numbers.py + tok_bins = [(64, 128), (128, 256), (256, 512)] + [ + (512 + 256 * i, 512 + 256 * (i + 1)) for i in range(14) + ] + + if args.only_first_n_bins is not None: + tok_bins = tok_bins[: args.only_first_n_bins] + + # Map token bins to pair-length bins using a simple heuristic + len_bins = [ + (lo // args.tokens_per_pair, hi // args.tokens_per_pair) + for (lo, hi) in tok_bins + ] + + # Optional dry-run: build a tiny sample to preview format and exit + if args.dry_run: + train, val, test = generate_kv_dataset( + n=3, + k=max(1, len_bins[0][0] or 1), + key_type=args.key_type, + value_type=args.value_type, + ) + print("Sample entry:") + print(json.dumps(train[0], indent=2)) + return + + for len_bin, tok_bin in zip(len_bins, tok_bins): + bin_size = max(1, len_bin[1] - len_bin[0]) # avoid div-by-zero for tiny bins + save_dir = f"{args.out_prefix}_{tok_bin[0]}_{tok_bin[1]}" + train_data, val_data, test_data = [], [], [] + + # Iterate over different context lengths (number of pairs) + for k in range(max(1, len_bin[0]), max(1, len_bin[1])): + # scale like the numbers script + per_k = max(1, args.base_samples_per_bin // bin_size) + train, val, test = generate_kv_dataset( + n=per_k, k=k, key_type=args.key_type, value_type=args.value_type + ) + train_data += train + val_data += val + test_data += test + + os.makedirs(save_dir, exist_ok=True) + 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 __name__ == "__main__": + main() diff --git a/data/generate_ctx_magic_number.py b/data/generate_ctx_magic_number.py new file mode 100644 index 0000000..5510c00 --- /dev/null +++ b/data/generate_ctx_magic_number.py @@ -0,0 +1,195 @@ +import argparse +import json +import math +import os +import random + +# ----------------------------- +# Config knobs (edit or use CLI) +# ----------------------------- +TOKENS_PER_BLOCK = 40 # rough heuristic tokens per noise block +BASE_SAMPLES_PER_BIN = 12_800 +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}." +SEP = "\n" # between blocks + + +def save_jsonl(data: list[dict], filepath: str) -> None: + parent_dir = os.path.dirname(filepath) + if parent_dir: + os.makedirs(parent_dir, exist_ok=True) + with open(filepath, "w") as f: + for entry in data: + json.dump(entry, f) + f.write("\n") + + +essential_digits4 = lambda: f"{random.randint(0, 9_999):04d}" + + +def _choose_position(total_blocks: int, depth_bin: int) -> int: + """Choose an insertion index for the special sentence within [0, total_blocks-1] + such that its relative depth falls within the depth bin [i/10, (i+1)/10). + """ + if total_blocks <= 0: + return 0 + # Use floor for start and ceil for end to cover boundaries evenly + start = math.floor(total_blocks * (depth_bin / 10)) + end = math.ceil(total_blocks * ((depth_bin + 1) / 10)) - 1 + # clamp + start = max(0, min(start, total_blocks - 1)) + end = max(start, min(end, total_blocks - 1)) + return random.randint(start, end) + + +def _build_example(total_blocks: int, depth_bin: int) -> dict: + """Build one example with a special line inserted among noise blocks. + + total_blocks: total number of blocks in the final context (including the special one) + depth_bin: integer in [0, 9] + """ + total_blocks = max(1, total_blocks) + + # Prepare blocks + magic = essential_digits4() + special_line = SPECIAL_TPL.format(magic_number=magic) + + # We'll have (total_blocks - 1) noise blocks and 1 special line + noise_count = max(0, total_blocks - 1) + blocks: list[str] = [NOISE_BLOCK for _ in range(noise_count)] + + insert_at = _choose_position(total_blocks, depth_bin) + # Insert special line at the desired position within the final sequence + # If noise_count == 0, we just return special + if noise_count == 0: + final_blocks = [special_line] + else: + # Compose by interleaving noise and inserting special at index + # Build a list of length `total_blocks` and fill + final_blocks = [] + noise_idx = 0 + for idx in range(total_blocks): + if idx == insert_at: + final_blocks.append(special_line) + else: + final_blocks.append(blocks[noise_idx]) + noise_idx += 1 + + context = SEP.join(final_blocks) + prompt = "What is the special magic number? Reply with only the number." + response = magic + 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 + base = n // 10 + rem = n % 10 + counts = [base + (1 if i < rem else 0) for i in range(10)] + + dataset: 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 + + +def main(): + parser = argparse.ArgumentParser( + 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( + "--base-samples-per-bin", + type=int, + default=BASE_SAMPLES_PER_BIN, + help="Baseline number of samples per token bin (scaled by bin width)", + ) + parser.add_argument( + "--out-prefix", + type=str, + default="data/raw_datasets/ctx_magic_number", + help="Output directory prefix (bin range will be appended)", + ) + parser.add_argument( + "--tokens-per-block", + "--tokens-per-pair", + dest="tokens_per_block", + type=int, + default=TOKENS_PER_BLOCK, + help="Heuristic tokens per noise block for bucketing", + ) + parser.add_argument( + "--only-first-n-bins", + type=int, + default=None, + help="For quick tests: only generate the first N token bins", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print a small sample and exit without writing files", + ) + + args = parser.parse_args() + + 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) + ] + + 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.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)) + return + + for len_bin, tok_bin in zip(len_bins, tok_bins): + bin_size = max(1, len_bin[1] - len_bin[0]) + save_dir = f"{args.out_prefix}_{tok_bin[0]}_{tok_bin[1]}" + train_data: list[dict] = [] + 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 + + os.makedirs(save_dir, exist_ok=True) + 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 __name__ == "__main__": + main() diff --git a/data/generate_ctx_numbers.py b/data/generate_ctx_numbers.py new file mode 100644 index 0000000..a74d959 --- /dev/null +++ b/data/generate_ctx_numbers.py @@ -0,0 +1,95 @@ +import json +import os +import random + + +def save_jsonl(data: list[dict], filepath: str) -> None: + """Save data to a JSONL file.""" + parent_dir = os.path.dirname(filepath) + if parent_dir: # Only create directories if there's a parent path + os.makedirs(parent_dir, exist_ok=True) + with open(filepath, "w") as f: + for entry in data: + json.dump(entry, f) + f.write("\n") + + +def get_random_combinations(numbers: list[int], n: int, k: int) -> list[list[int]]: + """Get n random combinations of k numbers from the list.""" + # using itertools.combinations hangs with large numbers + # so we explicitly generate the n indices + indices = [random.sample(range(len(numbers)), k=k) for _ in range(n)] + return [[numbers[i] for i in ind] for ind in indices] + + +def generate_number_dataset(max_num: int = 1000, n: int = 12000, k: int = 3): + """ + Generate a dataset of numbers with corresponding query and answer, + split into train/val/test sets. + + Args: + max_num: Maximum number in the range (exclusive) + k: Number of elements in each combination + """ + # Generate list of all numbers and shuffle them + numbers = list(range(max_num)) + random.shuffle(numbers) + + # Create dataset entries + dataset = [] + query = f"Repeat the information above exactly. Do not output anything else." + + # Generate all unique combinations of k numbers + combinations = get_random_combinations(numbers, n, k) + # random.shuffle(combinations) + + for combination in combinations: + num_list = [f"{i + 1}. {num}" for i, num in enumerate(combination)] + entry = { + "context": "\n".join(num_list), + "prompt": query, + "response": "\n".join(num_list), + } + dataset.append(entry) + + # Calculate split sizes + total_size = len(dataset) + train_size = int(0.98 * total_size) + val_size = int(0.01 * total_size) + + # Split dataset + train_data = dataset[:train_size] + val_data = dataset[train_size : train_size + val_size] + test_data = dataset[train_size + val_size :] + return train_data, val_data, test_data + + +if __name__ == "__main__": + # Set random seed for reproducibility + random.seed(42) + tok_bins = [(64, 128), (128, 256), (256, 512)] + [ + (512 + 256 * i, 512 + 256 * (i + 1)) for i in range(14) + ] + # roughly 9 tokens per number + tok_per_num = [9 if bin[0] >= 1024 else 8 for bin in tok_bins] + len_bins = [ + (bin[0] // tok, bin[1] // tok) for bin, tok in zip(tok_bins, tok_per_num) + ] + + for len_bin, tok_bin in zip(len_bins, tok_bins): + bin_size = len_bin[1] - len_bin[0] + save_dir = f"data/raw_datasets/ctx_numbers_{tok_bin[0]}_{tok_bin[1]}" + train_data, val_data, test_data = [], [], [] + for k in range(*len_bin): + train, val, test = generate_number_dataset(n=128_000 // bin_size, k=k) + train_data += train + val_data += val + test_data += test + + # Save splits to separate files + os.makedirs(save_dir, exist_ok=True) + 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}") diff --git a/data/self_generate_qa.py b/data/self_generate_qa.py index 19d9db3..717bac8 100644 --- a/data/self_generate_qa.py +++ b/data/self_generate_qa.py @@ -41,31 +41,21 @@ SYSTEM_TEMPLATE = ( "### END OF SYSTEM INSTRUCTION ###" ) +CTX_Q_SEP = "\n\n---\n\n# User\n" + SELF_QA_INTX = ( - "---\n\n" + # "---\n\n" "# System Instruction\n" "- The information provided is up-to-date information.\n" "- When the provided information is not relevant to the question, ***ignore*** it and answer the question based on your knowledge.\n" - "- If the provided information is related to the question, incorporate it in your response.\n" - "\n---\n\n" - "# User\n" + "- If the provided information is related to the question, incorporate it in your response." + + CTX_Q_SEP ) -PRE_CTX = "# Information\n" +PROMPT_TEMPLATE = "{context}" + CTX_Q_SEP + "{question}" -PROMPT_TEMPLATE = ( - PRE_CTX - + "{context}\n\n---\n\n" - + - # "# System Instruction\n" - # "- The information provided is up-to-date information.\n" - # "- When the provided information is not relevant to the question, ***ignore*** it and answer the question based on your knowledge.\n" - # "- If the provided information is related to the question, incorporate it in your response.\n" - # "\n---\n\n" - # "# User\n{question}" - SELF_QA_INTX - + "{question}" -) +PRE_CTX = "" +PROMPT_TEMPLATE_QA = "# Information\n{context}\n\n---\n\n" + SELF_QA_INTX + "{question}" """ @@ -114,8 +104,9 @@ def truncate_middle_if_too_long( return input_ids -def get_prompt(context: str, q: str) -> str: - return PROMPT_TEMPLATE.format(context=context, question=q) +def get_prompt(context: str, q: str, use_qa_template: bool) -> str: + template = PROMPT_TEMPLATE_QA if use_qa_template else PROMPT_TEMPLATE + return template.format(context=context, question=q) def add_closed_qa_prompt(q: str, closed_qa_prob: float = 0.1) -> str: @@ -187,7 +178,11 @@ def get_dataset_configs( def create_messages( - ctxs: list[str], questions: list[list[str]], vllm_model: str, system_template: str + ctxs: list[str], + questions: list[list[str]], + vllm_model: str, + system_template: str, + use_qa_template: bool, ) -> list[list[dict]]: """Create chat messages for the model.""" # if "gemma" in vllm_model: @@ -196,7 +191,9 @@ def create_messages( [ { "role": "user", - "content": system_template + "\n\n\n" + get_prompt(ctx, q), + "content": system_template + + "\n\n\n" + + get_prompt(ctx, q, use_qa_template), } ] for ctx, q_list in zip(ctxs, questions) @@ -221,6 +218,8 @@ def self_generate( system_template: str, parquet_file: str | None = None, do_truncate: bool = False, + use_qa_template: bool = True, + max_new_tokens: int = 1024, ) -> None: """Process a single dataset and generate QA pairs.""" @@ -254,6 +253,7 @@ def self_generate( print(f"Processing dataset: {ds_name}, split: {split}") print(f"Using temperature: {temp}") print(f"Using closed QA prompt probability: {closed_qa_prob}") + print(f"Use QA template: {use_qa_template}") if parquet_file: print(f"Loading dataset from parquet file: {parquet_file}") @@ -282,8 +282,8 @@ def self_generate( tk = get_tokenizer(args.vllm_model, train=True) - self_qa_intx_tokens = tk(SELF_QA_INTX, add_special_tokens=False)["input_ids"] - n_self_qa_intx_tokens = len(self_qa_intx_tokens) + ctx_q_sep_tokens = tk(CTX_Q_SEP, add_special_tokens=False)["input_ids"] + n_qa_sep_tokens = len(ctx_q_sep_tokens) pre_ctx_tokens = tk(PRE_CTX, add_special_tokens=False)["input_ids"] n_pre_ctx_tokens = len(pre_ctx_tokens) sys_tokens = tk(system_template.split("\n")[0], add_special_tokens=False)[ @@ -316,7 +316,11 @@ def self_generate( chunk_ctxs = ctxs[start : start + chunk_size] chunk_questions = questions[start : start + chunk_size] chunk_messages = create_messages( - chunk_ctxs, chunk_questions, args.vllm_model, SELF_GEN_SYSTEM_MSG + chunk_ctxs, + chunk_questions, + args.vllm_model, + SELF_GEN_SYSTEM_MSG, + use_qa_template, ) if do_truncate: @@ -330,7 +334,7 @@ def self_generate( truncate_middle_if_too_long( ids, max_length=MODEL_CTX_LEN[args.vllm_model], - max_new_tokens=1024, + max_new_tokens=max_new_tokens, ) for ids in tokenized_contents["input_ids"] ] @@ -350,8 +354,8 @@ def self_generate( llm, temp, tk, - self_qa_intx_tokens, - n_self_qa_intx_tokens, + ctx_q_sep_tokens, + n_qa_sep_tokens, sys_tokens, n_sys_tokens, chunk_ctxs, @@ -359,6 +363,7 @@ def self_generate( chunk_questions, chunk_messages, k, + max_new_tokens, ) @@ -368,8 +373,8 @@ def execute_qa_generation( llm, temp, tk, - self_qa_intx_tokens, - n_self_qa_intx_tokens, + ctx_q_sep_tokens, + n_qa_sep_tokens, sys_tokens, n_sys_tokens, ctxs, @@ -377,11 +382,12 @@ def execute_qa_generation( questions, messages, k, + max_tokens, ): completions = llm.chat( messages, sampling_params=SamplingParams( - max_tokens=1024, + max_tokens=max_tokens, logprobs=k, temperature=temp, seed=42, @@ -389,6 +395,8 @@ def execute_qa_generation( skip_special_tokens=False, include_stop_str_in_output=True, ), + chat_template=tk.chat_template, + chat_template_content_format="string", ) self_gen_data = { @@ -451,13 +459,13 @@ def execute_qa_generation( q_start = None for ii in range( - len(prompt_ids) - n_self_qa_intx_tokens, + len(prompt_ids) - n_qa_sep_tokens, -1, -1, ): - if prompt_ids[ii : ii + n_self_qa_intx_tokens] == self_qa_intx_tokens: + if prompt_ids[ii : ii + n_qa_sep_tokens] == ctx_q_sep_tokens: # found the start of the user input - q_start = ii + n_self_qa_intx_tokens + q_start = ii + n_qa_sep_tokens break # bos + question + eos + start model turn + response + eos @@ -611,6 +619,17 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Truncate contexts to fit model context length", ) + parser.add_argument( + "--max_new_tokens", + type=int, + default=1024, + help="Maximum number of new tokens to generate (default: 1024)", + ) + parser.add_argument( + "--remove_qa_template", + action="store_true", + help="Remove the QA template from the prompt (default: False)", + ) return parser.parse_args() @@ -650,7 +669,17 @@ if __name__ == "__main__": # Process each dataset for ds_name, split in dataset_configs: print(f"Processing dataset: {ds_name}, split: {split}") - self_generate(ds_name, split, args, llm, SYSTEM_TEMPLATE, args.do_truncate) + self_generate( + ds_name, + split, + args, + llm, + system_template=SYSTEM_TEMPLATE, + parquet_file=None, + do_truncate=args.do_truncate, + use_qa_template=not args.remove_qa_template, + max_new_tokens=args.max_new_tokens, + ) else: assert args.glob_pattern, ( "glob_pattern must be provided if no ds_names or config" @@ -666,4 +695,6 @@ if __name__ == "__main__": llm=llm, system_template=SYSTEM_TEMPLATE, do_truncate=args.do_truncate, + use_qa_template=not args.remove_qa_template, + max_new_tokens=args.max_new_tokens, ) diff --git a/pyproject.toml b/pyproject.toml index 6cdec56..0865184 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ "kaggle>=1.7.4.5", "bitsandbytes>=0.46.1", "google-cloud-storage>=3.2.0", + "wonderwords>=2.2.0", ] [build-system] diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index db7d26c..6513706 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -380,6 +380,13 @@ 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."}, + ) max_ctx_chunk_len: int = field( default=-1, metadata={ @@ -387,6 +394,12 @@ class CtxTrainingArguments: "it will be split up into multiple chunks." }, ) + min_ctx_chunk_len: int = field( + default=-1, + metadata={ + "help": "Min context chunk length. Used only with random chunking training" + }, + ) max_ctx_chunk_num: int = field( default=-1, metadata={"help": "Max number of context chunks per sample."}, diff --git a/src/ctx_to_lora/data/collator.py b/src/ctx_to_lora/data/collator.py index 9cfe554..69d9c0e 100644 --- a/src/ctx_to_lora/data/collator.py +++ b/src/ctx_to_lora/data/collator.py @@ -18,8 +18,10 @@ def flatten_if_not_packed(inp_list): if "position_ids" in sample: if n == 1: n_queries = sample.pop("n_queries") + n_ctx_chunks = sample.pop("n_ctx_chunks") batch = default_data_collator(inp_list, return_tensors="pt") batch["n_queries"] = torch.tensor(n_queries) + batch["n_ctx_chunks"] = torch.tensor(n_ctx_chunks) return batch elif n > 1: raise NotImplementedError("Please use batch_size=1 when using packed data") @@ -31,20 +33,29 @@ def flatten_if_not_packed(inp_list): assert not need_flatten, f"Validation data should not be nested." n_queries = torch.ones(len(inp_list), dtype=torch.int32) + n_ctx_chunks = torch.tensor( + [len(example["ctx_ids"]) for example in inp_list], dtype=torch.int32 + ) packed_inputs = flattener(inp_list, return_tensors="pt") packed_inputs["n_queries"] = n_queries + packed_inputs["n_ctx_chunks"] = n_ctx_chunks + if "ctx_ids" in sample: # HACK: assumes 1 ctx chunk here # TODO: fix this after we have chunk ctx training - assert all(len(ctx_ids) == 1 for ctx_ids in sample["ctx_ids"]), ( - "ctx_ids can only have one chunk for eval. " - "Please implement chunked ctx forward pass to handle this." - ) - ctx_ids = [{"input_ids": example["ctx_ids"][0]} for example in inp_list] - packed_ctx = flattener(ctx_ids, return_tensors="pt") - packed_inputs["ctx_ids"] = packed_ctx["input_ids"] - packed_inputs["ctx_position_ids"] = packed_ctx["position_ids"] + # print(sample["ctx_ids"]) + # breakpoint() + # assert all(len(ctx_ids) == 1 for ctx_ids in sample["ctx_ids"]), ( + # "ctx_ids can only have one chunk for eval. " + # "Please implement chunked ctx forward pass to handle this." + # ) + ctx_ids = concat_list([example.pop("ctx_ids") for example in inp_list]) + ctx_position_ids = torch.cat([torch.arange(len(ids)) for ids in ctx_ids]) + ctx_ids = torch.tensor(concat_list(ctx_ids)) + + packed_inputs["ctx_ids"] = ctx_ids.unsqueeze(0) + packed_inputs["ctx_position_ids"] = ctx_position_ids.unsqueeze(0) # for eval info if "ctx_ids_len" in sample: packed_inputs["ctx_ids_len"] = [ diff --git a/src/ctx_to_lora/data/definitions.py b/src/ctx_to_lora/data/definitions.py index 9433b8c..0b911fe 100644 --- a/src/ctx_to_lora/data/definitions.py +++ b/src/ctx_to_lora/data/definitions.py @@ -583,6 +583,30 @@ DS_KWARGS = { ), } +# add ctx_numbers +tok_bins = [(64, 128), (128, 256), (256, 512)] + [ + (512 + 256 * i, 512 + 256 * (i + 1)) for i in range(14) +] +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( + train=dict( + path="json", + data_files=f"data/raw_datasets/{toy_ds_name}_{tok_bin[0]}_{tok_bin[1]}/train.jsonl", + split="train", + ), + validation=dict( + path="json", + data_files=f"data/raw_datasets/{toy_ds_name}_{tok_bin[0]}_{tok_bin[1]}/val.jsonl", + split="train", + ), + test=dict( + path="json", + data_files=f"data/raw_datasets/{toy_ds_name}_{tok_bin[0]}_{tok_bin[1]}/test.jsonl", + split="train", + ), + ) + # LongBench kwargs for ds_name in LONGBENCH_TASKS + LONGBENCH_E_TASKS: DS_KWARGS[ds_name] = dict( diff --git a/src/ctx_to_lora/data/packing.py b/src/ctx_to_lora/data/packing.py index cbe0349..70ceed5 100644 --- a/src/ctx_to_lora/data/packing.py +++ b/src/ctx_to_lora/data/packing.py @@ -14,8 +14,7 @@ logger = logging.getLogger() def pack_data_points_by_length( lens: list[list[int]], - ctx_lens: list[int], - n_queries: list[int], + ctx_lens: list[list[int]], max_packed_inp_len: int, max_packed_ctx_len: int, max_size: int = -1, @@ -24,10 +23,9 @@ def pack_data_points_by_length( return [] len_arr = np.array([sum(l) for l in lens], dtype=np.long) - ctx_len_arr = np.array(ctx_lens, dtype=np.long) + ctx_len_arr = np.array([sum(l) for l in ctx_lens], dtype=np.long) n = len(len_arr) assert len(ctx_len_arr) == n, "Length of ctx_len_arr must match length of lens" - assert len(n_queries) == n, "Length of n_queries must match length of lens" if n == 1: return ( @@ -40,11 +38,6 @@ def pack_data_points_by_length( cumsum_inp_len = np.cumsum(len_arr) cumsum_ctx_len = np.cumsum(ctx_len_arr) - # indices of flattened input_ids - cumsum_n_queries = np.cumsum(n_queries) - - # start_indices_inp = [0] - # start_indices_ctx = [0] idx_pairs = [] i = 0 @@ -85,7 +78,7 @@ def pack_data_points_FA( raise ValueError("Batch is empty") # Pre-allocate lists with known sizes - total_ctx_len = sum(len(x) for x in batch["ctx_ids"]) + total_ctx_len = sum(len(y) for x in batch["ctx_ids"] for y in x) total_inp_len = sum(len(y) for x in batch["input_ids"] for y in x) ctx_ids = np.empty(total_ctx_len, dtype=np.long) @@ -111,7 +104,7 @@ def pack_data_points_FA( else: sequences = zip(batch["input_ids"], batch["labels"]) - offset = ctx_offset = 0 + offset = 0 for sample in sequences: input_ids_b, labels_b = sample[:2] inp_start = offset @@ -119,10 +112,9 @@ def pack_data_points_FA( # compute position_ids for each sub-list in input_ids_b local_start = inp_start for ids_b in input_ids_b: - position_ids[local_start : local_start + len(ids_b)] = np.arange( - len(ids_b), dtype=np.long - ) - local_start += len(ids_b) + local_end = local_start + len(ids_b) + position_ids[local_start:local_end] = np.arange(len(ids_b), dtype=np.int32) + local_start = local_end input_ids_b = concat_list(input_ids_b) labels_b = concat_list(labels_b) @@ -145,11 +137,20 @@ def pack_data_points_FA( ) logits_offset += logits_len + ctx_offset = 0 for ctx_ids_b in batch["ctx_ids"]: + local_start = ctx_offset + for ctx_ids_b_item in ctx_ids_b: + local_end = local_start + len(ctx_ids_b_item) + ctx_position_ids[local_start:local_end] = np.arange( + len(ctx_ids_b_item), dtype=np.int32 + ) + local_start = local_end + + ctx_ids_b = concat_list(ctx_ids_b) ctx_len = len(ctx_ids_b) ctx_start, ctx_end = ctx_offset, ctx_offset + ctx_len ctx_ids[ctx_start:ctx_end] = ctx_ids_b - ctx_position_ids[ctx_start:ctx_end] = np.arange(ctx_len, dtype=np.long) ctx_offset += ctx_len out = { @@ -179,18 +180,18 @@ def pack_batch( ) n_queries = [len(x) for x in batch["input_ids"]] + n_ctx_chunks = [len(x) for x in batch["ctx_ids"]] inp_lens = [[len(y) for y in x] for x in batch["input_ids"]] inp_count = len(inp_lens) if "ctx_ids" not in batch: raise ValueError("Batch must contain 'ctx_ids' and 'labels' keys") # we do not pad so we can just take the length of the tokens - ctx_lens = [len(x) for x in batch["ctx_ids"]] + ctx_lens = [[len(y) for y in x] for x in batch["ctx_ids"]] # Group indices idx_pairs = pack_data_points_by_length( inp_lens, ctx_lens, - n_queries, max_packed_inp_len, max_packed_ctx_len, max_packed_size, @@ -205,6 +206,7 @@ def pack_batch( "position_ids": [], "labels": [], "n_queries": [], + "n_ctx_chunks": [], } has_logprobs = "logprobs_vals" in batch if has_logprobs: @@ -235,6 +237,7 @@ def pack_batch( packed_batch["position_ids"].append(packed_item["position_ids"]) packed_batch["labels"].append(packed_item["labels"]) packed_batch["n_queries"].append(n_queries[start_idx:end_idx]) + packed_batch["n_ctx_chunks"].append(n_ctx_chunks[start_idx:end_idx]) if has_logprobs: packed_batch["logprobs_vals"].append(packed_item["logprobs_vals"]) packed_batch["logprobs_indices"].append(packed_item["logprobs_indices"]) diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index 5b5c23a..ba58dcf 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -10,6 +10,7 @@ from os import path from typing import Any import datasets +import numpy as np import torch from datasets import Dataset, interleave_datasets, is_caching_enabled, load_dataset from transformers import PreTrainedTokenizerBase @@ -210,6 +211,8 @@ def get_tokenized_dataset( ctx_model_max_len: int, ctx_tokenizer: PreTrainedTokenizerBase, max_ctx_chunk_len: int, + min_ctx_chunk_len: int, + random_chunking: bool, max_ctx_chunk_num: int, add_ctx_to_chat: bool, use_kl_loss: bool, @@ -234,6 +237,8 @@ def get_tokenized_dataset( ctx_model_max_len=ctx_model_max_len, 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, max_ctx_chunk_num=max_ctx_chunk_num, need_ctx_ids=need_ctx_ids, split=split, @@ -304,6 +309,8 @@ def construct_and_tokenize_ctx_qa( add_ctx_to_chat, need_ctx_ids, max_ctx_chunk_len, + min_ctx_chunk_len, + random_chunking, max_ctx_chunk_num, ds, split, @@ -354,14 +361,14 @@ def construct_and_tokenize_ctx_qa( remove_columns=["context"], ) - # TODO: this can be removed once we implement ctx chunking for training - if is_train: - # drop - tokenized_ds = tokenized_ds.filter( - len_filter, - fn_kwargs={"max_length": ctx_model_max_len, "keys": ["ctx_ids"]}, - num_proc=16, - ) + # # TODO: this can be removed once we implement ctx chunking for training + # if is_train: + # # drop + # tokenized_ds = tokenized_ds.filter( + # len_filter, + # fn_kwargs={"max_length": ctx_model_max_len, "keys": ["ctx_ids"]}, + # num_proc=16, + # ) # ctx chunking # Ideally we want to chunk the raw text directly @@ -369,23 +376,30 @@ def construct_and_tokenize_ctx_qa( # we can only chunk the tokenized context which requires some workaround # e.g., removing/applying template to each chunk # with some big caveats, e.g., losing order info - tokenized_ds = tokenized_ds.map( - split_too_long_ctx, - fn_kwargs={ - "max_length": max_ctx_chunk_len, - "max_num_split": max_ctx_chunk_num, - "model_name_or_path": tokenizer.name_or_path, - "is_train": is_train, - }, + 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, + "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) + logging.info( + f"Avg. num chunks per ctx: {np.mean(list(map(len, tokenized_ds['ctx_ids'])))}" ) - logging.info(f"Split too long QAs with max length {max_qas_len}") + split_qa_kwargs = { + "max_qas_len": max_qas_len, + "max_qas_per_sample": max_qas_per_sample, + } + logging.info(f"Split too long QAs with {split_qa_kwargs=}") tokenized_ds = tokenized_ds.map( split_too_long_qas, - fn_kwargs={ - "max_qas_len": max_qas_len, - "max_qas_per_sample": max_qas_per_sample, - }, + fn_kwargs=split_qa_kwargs, batched=True, batch_size=12_500, num_proc=16, @@ -562,57 +576,56 @@ def convert_ctx_prompt_response_to_messages( def split_too_long_ctx( sample: dict[str, Any], - max_length: int, - max_num_split: int, model_name_or_path: str, + max_chunk_len: int, + min_chunk_len: int, + max_num_split: int, is_train: bool, + random_chunking: bool, ) -> dict[str, Any]: """ Split context into smaller chunks if it exceeds the maximum length. Args: samples: Dictionary containing 'ctx_ids' and 'ctx_attn_mask' - max_length: Maximum length for each context chunk + 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 length since merging - # a batch of chunked loras need padded in the rank axes + # 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) - # - # TODO: should be stochastic during training? + # 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 - # For eval, always split at max len? - # raise NotImplementedError( - # "Splitting contexts for training is not implemented yet. " - # ) - # # Limit the number of splits for training? - # if len(chunks) > max_num_split > 0: - # chunks = chunks[:max_num_split] - # TODO: training inputs right now is not nested - return sample + if random_chunking: + chunk_len = random.randint(min_chunk_len, max_chunk_len) ctx_affixes = CTX_AFFIXES[model_name_or_path] prefix = ctx_affixes["prefix"] suffix = ctx_affixes["suffix"] ctx_ids = sample["ctx_ids"] - if max_length <= 0 and max_num_split <= 0: + if chunk_len <= 0 and max_num_split <= 0: return {"ctx_ids": [ctx_ids]} - if len(ctx_ids) <= max_length: + if len(ctx_ids) <= chunk_len: return {"ctx_ids": [ctx_ids]} - # can we improve how we chunk? - # compute num chunks for more uniform chunking - n_chunks = ceil(len(ctx_ids) / max_length) + # uniform chunking + n_chunks = ceil(len(ctx_ids) / chunk_len) avg_len = ceil(len(ctx_ids) / n_chunks) # Split the context into smaller chunks @@ -815,9 +828,14 @@ def add_length_info(sample: dict[str, any], columns: list[str]) -> dict[str, any out = {} for k in columns: if check_is_iterable(sample[k][0]): + # ctx_ids out[f"{k}_len"] = sum([len(x) for x in sample[k]]) else: - out[f"{k}_len"] = len(sample[k]) + # input_ids + label_idx = None + if k == "input_ids" and "labels" in sample: + label_idx = np.argmax(np.array(sample["labels"]) != -100) + out[f"{k}_len"] = len(sample[k][:label_idx]) return out diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py index 9d8afd1..a3e3254 100644 --- a/src/ctx_to_lora/eval_utils.py +++ b/src/ctx_to_lora/eval_utils.py @@ -788,6 +788,8 @@ def evaluate( max_qas_len=-1, 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, base_model_max_len=model.base_model.config.max_position_embeddings, tokenizer=tokenizer, diff --git a/src/ctx_to_lora/metrics.py b/src/ctx_to_lora/metrics.py index ecd2679..9e33931 100644 --- a/src/ctx_to_lora/metrics.py +++ b/src/ctx_to_lora/metrics.py @@ -75,14 +75,15 @@ def compute_prefix_matching(shift_logits, shift_labels, valid_masks): @torch.inference_mode() def compute_perplexity(shift_logits, shift_labels, valid_masks): - loss_fct = torch.nn.CrossEntropyLoss(reduction="none") - loss = loss_fct(shift_logits.transpose(1, 2), shift_labels) - loss = (loss * valid_masks).sum(dim=1) / valid_masks.sum(dim=1) - preplexities = torch.exp(loss) - return { - "perplexities": preplexities.tolist(), - "n_perplexities": valid_masks.shape[0], - } + return {"perplexities_ph": [1], "n_perplexities_ph": 1} + # loss_fct = torch.nn.CrossEntropyLoss(reduction="none") + # loss = loss_fct(shift_logits.transpose(1, 2), shift_labels) + # loss = (loss * valid_masks).sum(dim=1) / valid_masks.sum(dim=1) + # preplexities = torch.exp(loss) + # return { + # "perplexities": preplexities.tolist(), + # "n_perplexities": valid_masks.shape[0], + # } class Evaluator: diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index b592cc9..0aa9e1b 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/src/ctx_to_lora/modeling/hypernet.py @@ -226,12 +226,14 @@ class HyperLoRA(nn.Module): self._init_model() def _init_model(self): + # TODO: add a way to handle dependencies between chunks self.agg_config = self.config.aggregator_config self.aggregator = AGGREGATOR_CLS[self.agg_config.aggregator_type]( **vars(self.agg_config) ) self.lora_config = self.config.lora_config + self.r = self.lora_config.r self.target_modules = ( tuple(sorted(self.lora_config.target_modules)) if self.lora_config else None @@ -273,6 +275,23 @@ class HyperLoRA(nn.Module): 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 + # } + # ) + if self.config.use_light_weight_lora: # light-weight lora projection (per layer, per module) # self.pre_lora_projection = nn.ParameterDict( @@ -527,25 +546,17 @@ class HyperLoRA(nn.Module): # hidden_size=self.config.base_hidden_size, # ) - def get_head_bias(self): - bias_list = unpack( - self.head.bias, - [[] for _ in range(len(self.target_modules))], - "bs n_layers * r max_io_dim", - ) - bias_dict = dict() - for module, bias in zip(self.target_modules, bias_list): - bias_A, bias_B = unpack( - bias[..., : self.d_in[module] + self.d_out[module]], - [[self.d_in[module]], [self.d_out[module]]], - "bs n_layers r *", - ) - # 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 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, ...] + + # # 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 - # @torch.autocast(device_type="cuda", dtype=torch.bfloat16) def _to_lora_dict( self, flat_loras: Float[Tensor, "bs n_layers n_modules r max_io_dim"] ) -> dict[str, dict[str, Float[Tensor, "bs n_layers r _"]]]: @@ -651,6 +662,7 @@ 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__() @@ -660,6 +672,7 @@ 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 = [] @@ -773,8 +786,9 @@ class ModulatedPretrainedModel(nn.Module): 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: - A = A[0:1] - B = B[:, 0:1] + if not self.hypernet.config.use_per_rank_bias: + A = A[0:1] + B = B[:, 0:1] biases = [A, B.T] # bias-hyperinit @@ -878,6 +892,7 @@ class ModulatedPretrainedModel(nn.Module): ctx_ids: Integer[Tensor, "n_ctx ctx_len"] | None = None, ctx_attn_mask: Integer[Tensor, "n_ctx ctx_len"] | None = None, ctx_position_ids: Integer[Tensor, "n_ctx ctx_len"] | None = None, + n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None, n_queries: Integer[Tensor, "n_ctx"] | None = None, return_generated_lora: bool | None = False, *model_inputs_args: Any, @@ -917,6 +932,13 @@ class ModulatedPretrainedModel(nn.Module): ctx_ids, ctx_attn_mask, ctx_position_ids ) + generated_loras = combine_lora( + generated_loras, + n_ctx_chunks, + aggregation=self.lora_aggregation, + # lora_bias=self.hypernet.get_head_bias(), + ) + # input_ids in model_inputs_kwargs contains only # prompt + response (for hypernet training) position_ids = ( @@ -988,7 +1010,7 @@ class ModulatedPretrainedModel(nn.Module): 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_merger.py b/src/ctx_to_lora/modeling/lora_merger.py index 7cc74dc..e51b8f1 100644 --- a/src/ctx_to_lora/modeling/lora_merger.py +++ b/src/ctx_to_lora/modeling/lora_merger.py @@ -10,18 +10,8 @@ from jaxtyping import Integer from torch import Tensor -def compute_rank(n_lora, rank, aggregation): - """Return the effective (max) rank required for a group. - - aggregation == 'mean': we only concatenate the existing LoRA ranks (n_lora * r) - aggregation == 'sum' : an additional bias term may be appended (+1 * r) - """ - if aggregation == "mean": - return n_lora * rank - elif aggregation == "sum": - return (n_lora + 1) * rank - else: - raise ValueError(f"Unknown aggregation: {aggregation}") +def compute_rank(n_lora, rank): + return (n_lora + 1) * rank def combine_lora( @@ -53,22 +43,20 @@ def combine_lora( combined_A shape: [G, n_layers, max_rank, dim] combined_B shape: [G, n_layers, dim, max_rank] """ - assert aggregation in ["mean", "sum"] + # TODO: optimize! + # assert aggregation in ["mean", "sum"] + if (n_chunks == 1).all(): + return generated_loras # Precompute common quantities - num_groups = len(n_chunks) total_chunks = int(n_chunks.sum()) - sqrt_group_sizes = n_chunks**0.5 # [G] - # Repeat for each chunk within a group: e.g., - # if n_chunks=[2,3] -> repeat factors [sqrt2, sqrt2, sqrt3, sqrt3, sqrt3] - sqrt_group_sizes_per_chunk = sqrt_group_sizes.repeat_interleave(n_chunks, dim=0) # 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, aggregation) for n in n_chunks) + max_rank_needed = max(compute_rank(n, base_rank) for n in n_chunks) # Initialize output container combined_loras: dict[str, dict[str, Tensor]] = { @@ -84,15 +72,14 @@ def combine_lora( # or [total_chunks, n_layers, dim, r] (B) loras = module_loras[matrix_key] - # Expand bias if provided - bias_tensor = None - if lora_bias is not None: - bias_tensor = lora_bias[module_name][matrix_key].expand( - num_groups, *loras.shape[1:] - ) - # 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) # Split the leading chunk dimension into groups per n_chunks specification @@ -108,6 +95,14 @@ def combine_lora( 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, @@ -165,12 +160,10 @@ def _combine_single_matrix( 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}" + # assert matrix_key in ("A", "B"), f"matrix_key must be 'A' or 'B', got {matrix_key}" num_groups = len(n_chunks) - combined_shape = list(bias.shape) - combined_shape[0] = num_groups - + 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) @@ -178,51 +171,19 @@ def _combine_single_matrix( rank_dim_index = -1 # (G, layers, dim, rank) combined_shape[rank_dim_index] = max_rank_needed - combined = torch.zeros(*combined_shape, device=bias.device, dtype=bias.dtype) + 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): - # we remove leading dim when rearrange combined_rank = deltas.shape[rank_dim] - # Build slice pattern: rank slice up to combined_rank. + # Build slice pattern, slice up to combined_rank. slice_pattern = [g, slice(None), slice(None), slice(None)] slice_pattern[rank_dim] = slice(combined_rank) - if bias is not None and aggregation == "sum": - if matrix_key == "A": - deltas = rearrange( - deltas, - "1 n_layers (chunks r) dim -> chunks n_layers r dim", - chunks=n_chunks[g], - ) - else: - deltas = rearrange( - deltas, - "1 n_layers dim (chunks r) -> chunks n_layers dim r", - chunks=n_chunks[g], - ) - center = deltas.mean(dim=0, keepdim=True) - deltas = deltas - center - - if matrix_key == "A": - deltas = rearrange( - deltas, "chunks n_layers r dim -> 1 n_layers (chunks r) dim" - ) - cat_lora = torch.cat( - [deltas, center], - dim=rank_dim, - ) - combined[g, :, : (1 + n_chunks[g]) * base_rank, :] = cat_lora - else: - deltas = rearrange( - deltas, "chunks n_layers dim r -> 1 n_layers dim (chunks r)" - ) - combined[g, :, :, : (1 + n_chunks[g]) * base_rank] = torch.cat( - [deltas, center], - dim=rank_dim, - ) - - else: - combined[slice_pattern] = deltas + combined[slice_pattern] = deltas return combined diff --git a/src/ctx_to_lora/trainer.py b/src/ctx_to_lora/trainer.py index 7685502..6d844c4 100644 --- a/src/ctx_to_lora/trainer.py +++ b/src/ctx_to_lora/trainer.py @@ -302,16 +302,16 @@ def train_model( logger.info("Training with modulated model. Using ModulatedModelTrainer.") trainer_cls = ModulatedModelTrainer trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef - trainer_kwargs["use_per_ctx_average_loss"] = ( - training_args.use_per_ctx_average_loss - ) del training_args.gen_lora_l1_reg_coef - del training_args.use_per_ctx_average_loss if training_args.use_kl_loss: logger.info("Training with distillation loss. Using DistillationTrainer.") + trainer_kwargs["use_per_ctx_average_loss"] = ( + training_args.use_per_ctx_average_loss + ) trainer_cls = DistillationTrainer del training_args.use_kl_loss + del training_args.use_per_ctx_average_loss if training_args.auto_find_batch_size: # set the batch size to some high number @@ -319,7 +319,7 @@ def train_model( training_args.per_device_train_batch_size = 128 trainer = trainer_cls(**trainer_kwargs) - if trainer.use_per_ctx_average_loss: + if getattr(trainer, "use_per_ctx_average_loss", False): trainer.get_batch_samples = trainer.get_batch_samples_ctx # MONKEY PATCH: remove embedding layers from weight decay diff --git a/train.py b/train.py index 94a8850..5cd2cc2 100755 --- a/train.py +++ b/train.py @@ -199,6 +199,7 @@ def main(): base_model, hypernet_config, ctx_encoder_args, + lora_aggregation=ctx_args.lora_aggregation, ) else: @@ -249,6 +250,9 @@ def main(): ctx_model_max_len = model.ctx_encoder.config.max_position_embeddings if ctx_args.max_ctx_len > 0: ctx_model_max_len = ctx_args.max_ctx_len + if ctx_args.max_ctx_chunk_len <= 0: + # set default chunk size to max length of the ctx encoder + ctx_args.max_ctx_chunk_len = ctx_model_max_len _get_tokenized_dataset = partial( get_tokenized_dataset, @@ -260,6 +264,8 @@ def main(): ctx_tokenizer=ctx_tokenizer, 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, max_ctx_chunk_num=ctx_args.max_ctx_chunk_num, use_kl_loss=ctx_args.use_kl_loss, ) diff --git a/uv.lock b/uv.lock index 2245717..4701ac2 100644 --- a/uv.lock +++ b/uv.lock @@ -723,6 +723,7 @@ dependencies = [ { name = "transformers" }, { name = "vllm" }, { name = "wandb" }, + { name = "wonderwords" }, ] [package.metadata] @@ -757,6 +758,7 @@ requires-dist = [ { name = "transformers", specifier = "==4.51.3" }, { name = "vllm", specifier = "==0.8.5" }, { name = "wandb" }, + { name = "wonderwords", specifier = ">=2.2.0" }, ] [[package]] @@ -5787,6 +5789,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ca/51/5447876806d1088a0f8f71e16542bf350918128d0a69437df26047c8e46f/widgetsnbextension-4.0.14-py3-none-any.whl", hash = "sha256:4875a9eaf72fbf5079dc372a51a9f268fc38d46f767cbf85c43a36da5cb9b575", size = 2196503, upload-time = "2025-04-10T13:01:23.086Z" }, ] +[[package]] +name = "wonderwords" +version = "2.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f3/ed/d5f4336a7b1b416d5c502a6ed323f02acdb5e67467ef8f5e80ea6617fe3e/wonderwords-2.2.0.tar.gz", hash = "sha256:0b7ec6f591062afc55603bfea71463afbab06794b3064d9f7b04d0ce251a13d0", size = 50579, upload-time = "2021-02-17T20:28:45.402Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b4/a4/6a7705c6e17de71e45782a0284e06691f85350adb80b1dea9c027c07f024/wonderwords-2.2.0-py3-none-any.whl", hash = "sha256:65fc665f1f5590e98f6d9259414ea036bf1b6dd83e51aa6ba44473c99ca92da1", size = 44970, upload-time = "2021-02-17T20:28:43.977Z" }, +] + [[package]] name = "wrapt" version = "1.17.2"