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

toy ctx nums and multi-lora training (#11)

Browse source 
* multi-lora trainable toy number repeat dataset

* per rank bias init

* remove head_bias +simplify merge + skip perplexities metric

* ctx_numbers train example

* self-gen ctx numbers example
This commit is contained in:
Rujikorn Charakorn 2025-08-18 18:38:07 +09:00 • committed by GitHub
parent dfee685678
commit c5e9bc769d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
36 changed files with 1172 additions and 231 deletions
Download patch file Download diff file Expand all files Collapse all files

60
README.md
View file

@ -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
```
<!-- ### 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
```
@ -46,20 +55,29 @@ WANDB_MODE=disabled run uv run train.py configs/fw_qa_v2_level_0_tiny.yaml --mod
### 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
```
``` -->
### 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

2
chat_templates/google/gemma-2-9b-it.jinja
View file

@ -14,7 +14,7 @@
{%- if message['role'] == 'user' and loop.first and system_message is defined %}
{{ '<start_of_turn>' + role + '\n' + system_message + '\n\n' + message['content'] | trim + '<end_of_turn>\n' }}
{%- elif message['role'] == 'assistant' %}
{{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] | trim + '<end_of_turn>\n' }}{% endgeneration %}
{{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] + '<end_of_turn>' }}{% endgeneration %}{{ '\n' }}
{%- else %}
{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}
{%- endif %}

2
chat_templates/google/gemma-3-1b-it.jinja
View file

@ -22,7 +22,7 @@
{%- if (message['role'] == 'user') -%}
{{ '<start_of_turn>' + role + '\n' + (first_user_prefix if loop.first else "") + message['content'] | trim + '<end_of_turn>\n' }}
{%- elif (message['role'] == 'assistant') -%}
{{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] | trim + '<end_of_turn>\n' }}{% endgeneration %}
{{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] + '<end_of_turn>' }}{% endgeneration %}{{ '\n' }}
{%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}

2
chat_templates/google/gemma-3-4b-it.jinja
View file

@ -22,7 +22,7 @@
{%- if (message['role'] == 'user') -%}
{{ '<start_of_turn>' + role + '\n' + (first_user_prefix if loop.first else "") + message['content'] | trim + '<end_of_turn>\n' }}
{%- elif (message['role'] == 'assistant') -%}
{{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] | trim + '<end_of_turn>\n' }}{% endgeneration %}
{{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] + '<end_of_turn>' }}{% endgeneration %}{{ '\n' }}
{%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}

57
configs/gemma-3-1b-it/toy_exp/context_numbers_10.yaml Normal file
View file

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

12
configs/gemma-3-1b-it/toy_exp/ctx_kv_64_128.yaml Normal file
View file

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

12
configs/gemma-3-1b-it/toy_exp/ctx_kv_64_128_self_gen.yaml Normal file
View file

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

14
configs/gemma-3-1b-it/toy_exp/ctx_kv_64_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_kv_64_128
- ctx_kv_128_256
val_ds_names:
- ctx_kv_64_128
- ctx_kv_128_256

21
configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_1024.yaml Normal file
View file

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

12
configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_128.yaml Normal file
View file

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

12
configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_128_self_gen.yaml Normal file
View file

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

29
configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_2048.yaml Normal file
View file

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

14
configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_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_numbers_64_128
- ctx_numbers_128_256
val_ds_names:
- ctx_numbers_64_128
- ctx_numbers_128_256

17
configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_512.yaml Normal file
View file

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

20
configs/toy_exp/ctx_kv_64_1024_self_gen.yaml Normal file
View file

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

12
configs/toy_exp/ctx_kv_64_128_self_gen.yaml Normal file
View file

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

14
configs/toy_exp/ctx_kv_64_256_self_gen.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:
- 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

12
configs/toy_exp/ctx_numbers_64_128_self_gen.yaml Normal file
View file

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

28
configs/toy_exp/ctx_numbers_64_2048_self_gen.yaml Normal file
View file

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

243
data/generate_ctx_kv.py Normal file
View file

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

195
data/generate_ctx_magic_number.py Normal file
View file

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

95
data/generate_ctx_numbers.py Normal file
View file

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

101
data/self_generate_qa.py
View file

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

1
pyproject.toml
View file

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

13
src/ctx_to_lora/configs.py
View file

@ -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."},

27
src/ctx_to_lora/data/collator.py
View file

@ -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"] = [

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

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

39
src/ctx_to_lora/data/packing.py
View file

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

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

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

2
src/ctx_to_lora/eval_utils.py
View file

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

17
src/ctx_to_lora/metrics.py
View file

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

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

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

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

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

10
src/ctx_to_lora/trainer.py
View file

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

6
train.py
View file

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

11
uv.lock generated
View file

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