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

Refactor_and_improve_data (#7)

Browse source 
* faster slice

* add facts + ctx_qa

* new configs

* new scripts

* intx_sft.py to train.py

* add kaggle for downloading facts

* max_new_tokens cli for eval

* generate negative_nq

* scripts + configs

* default vals

* max_val_samples_per_ds=500

* more efficient layer-to-layer ctx encoder

* use_per_ctx_average_loss

* faster processing

* small exp distill

* scripts

* more robust watcher

* per-module l1_norm avg

* per-ctx average loss

* clear_gpu
This commit is contained in:
Rujikorn Charakorn 2025-08-04 20:52:37 +09:00 • committed by GitHub
parent 610750e6fb
commit 891c0bd256
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
75 changed files with 221183 additions and 2435 deletions
Download patch file Download diff file Expand all files Collapse all files

25
README.md
View file

@ -11,7 +11,7 @@
<!-- ### Finetuning the base model with LoRA adaptor
```bash
WANDB_MODE=disabled python hyperlora/intx_sft.py configs/default.yaml --model_name_or_path=meta-llama/Llama-3.2-1B-Instruct --num_train_epochs=50 --per_device_train_batch_size=128 --per_device_eval_batch_size=128 --exp_setup=lora --train_ds_name=data/raw_datasets/context_numbers_small
WANDB_MODE=disabled python hyperlora/train.py configs/default.yaml --model_name_or_path=meta-llama/Llama-3.2-1B-Instruct --num_train_epochs=50 --per_device_train_batch_size=128 --per_device_eval_batch_size=128 --exp_setup=lora --train_ds_name=data/raw_datasets/context_numbers_small
``` -->
## 🚀 API Usage [WIP]
@ -31,21 +31,21 @@ print(model.decode(outputs))
## 🏋️ Training
### 🔢 HyperLoRA w/ context_numbers_10
```bash
WANDB_MODE=disabled run uv run intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=1 --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=1000 --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=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208
WANDB_MODE=disabled run uv run train.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=1 --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=1000 --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=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208
```
KL loss
```bash
WANDB_MODE=disabled run uv run intx_sft.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
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
```bash
WANDB_MODE=disabled run uv run intx_sft.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
WANDB_MODE=disabled run uv run train.py configs/fw_qa_v2_level_0_tiny.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=64 --gradient_accumulation_steps=8 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=1 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --max_packed_inp_len=24000 --max_packed_ctx_len=48000 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=50
```
### Squad only
```bash
WANDB_MODE=disabled uv run intx_sft.py configs/squad.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=64 --gradient_accumulation_steps=8 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=1 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --max_packed_inp_len=16000 --max_packed_ctx_len=32000 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=10
WANDB_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
```
<!--
### HyperLoRA w/ self-gen 3 mini
@ -56,18 +56,18 @@ uv run python data/self_generate_qa.py \
--vllm_model=google/gemma-2-2b-it --config=configs/self_gen_3_mini.yaml
WANDB_MODE=disabled uv run python intx_sft.py configs/self_gen_3_mini.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=10 --per_device_train_batch_size=32 --gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=2 --lora_r=8 --eval_steps=5000 --save_steps=5000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1
WANDB_MODE=disabled uv run python train.py configs/self_gen_3_mini.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=10 --per_device_train_batch_size=32 --gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=2 --lora_r=8 --eval_steps=5000 --save_steps=5000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1
```
### HyperLoRA w/ fw-qa pretrain only
```bash
WANDB_MODE=disabled run uv run python intx_sft.py configs/fw_qa_pretrain_only.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=4 --gradient_accumulation_steps=8 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=4 --num_latent_factor=1 --lora_r=8 --eval_steps=100 --save_steps=100 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1
WANDB_MODE=disabled run uv run python train.py configs/fw_qa_pretrain_only.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=4 --gradient_accumulation_steps=8 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=4 --num_latent_factor=1 --lora_r=8 --eval_steps=100 --save_steps=100 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1
``` -->
<!-- ### HyperLoRA w/ context_numbers_128
```bash
WANDB_MODE=disabled run python hyperlora/intx_sft.py configs/context_numbers_128.yaml --model_name_or_path=meta-llama/Llama-3.2-1B-Instruct --num_train_epochs=10 --per_device_train_batch_size=64 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj
WANDB_MODE=disabled run python hyperlora/train.py configs/context_numbers_128.yaml --model_name_or_path=meta-llama/Llama-3.2-1B-Instruct --num_train_epochs=10 --per_device_train_batch_size=64 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj
``` -->
### Favourite numbers data generation
```bash
@ -79,6 +79,7 @@ uv run python data/generate_fav_num_big.py
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
```
### SQuAD
@ -173,22 +174,22 @@ uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --confi
<!--
### GSM8k LoRA
```bash
run python intx_sft.py configs/gsm8k.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=10 --per_device_train_batch_size=16 --gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=lora --target_modules=up_proj,down_proj --eval_steps=5000 --save_steps=5000 --learning_rate=1e-4 --neftune_noise_alpha=5 --load_best_model_at_end=True --metric_for_best_model=gsm8k_loss
run python train.py configs/gsm8k.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=10 --per_device_train_batch_size=16 --gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=lora --target_modules=up_proj,down_proj --eval_steps=5000 --save_steps=5000 --learning_rate=1e-4 --neftune_noise_alpha=5 --load_best_model_at_end=True --metric_for_best_model=gsm8k_loss
```
### Multitask LoRA
```bash
run python intx_sft.py configs/math_and_code.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=4 --gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=lora --target_modules=up_proj,down_proj --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --neftune_noise_alpha=5 --load_best_model_at_end=True --metric_for_best_model=gsm8k_loss
run python train.py configs/math_and_code.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=4 --gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=lora --target_modules=up_proj,down_proj --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --neftune_noise_alpha=5 --load_best_model_at_end=True --metric_for_best_model=gsm8k_loss
```
### HyperLoRA finetune on GSM8K
```bash
run python intx_sft.py configs/gsm8k.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=16 --gradient_accumulation_steps=1 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=1 --num_self_attends_per_block=16 --self_attention_widening_factor=1 --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --neftune_noise_alpha=5 --use_light_weight_lora=True --light_weight_latent_size=512 --load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False --add_repeat_prompt=False --ctx_encoder_model_name_or_path=meta-llama/Llama-3.2-11B-Vision-Instruct
run python train.py configs/gsm8k.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=16 --gradient_accumulation_steps=1 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=1 --num_self_attends_per_block=16 --self_attention_widening_factor=1 --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --neftune_noise_alpha=5 --use_light_weight_lora=True --light_weight_latent_size=512 --load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False --add_repeat_prompt=False --ctx_encoder_model_name_or_path=meta-llama/Llama-3.2-11B-Vision-Instruct
``` -->
### Continue from a checkpoint
```bash
run python intx_sft.py configs/...yaml ... --from_pretrained_checkpoint=train_outputs/runs/May09_16-25-35_slurm0-a3nodeset-4_59459_ea85a571/checkpoint-10000/pytorch_model.bin --resume_from_checkpoint=train_outputs/runs/May09_16-25-35_slurm0-a3nodeset-4_59459_ea85a571/checkpoint-10000
run python train.py configs/...yaml ... --from_pretrained_checkpoint=train_outputs/runs/May09_16-25-35_slurm0-a3nodeset-4_59459_ea85a571/checkpoint-10000/pytorch_model.bin --resume_from_checkpoint=train_outputs/runs/May09_16-25-35_slurm0-a3nodeset-4_59459_ea85a571/checkpoint-10000
```
### Evaluation

33
configs/context_numbers_10_self_gen.yaml
View file

@ -1,34 +1,5 @@
output_dir: "" # just a placeholder
bf16: true
model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
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_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
@ -38,4 +9,4 @@ train_ds_names:
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/context_numbers_2_10
val_ds_names:
- context_numbers_2_10
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/context_numbers_2_10

53
configs/small_exp/qa_short_ctx_lv1_small.yaml
View file

@ -1,53 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
# data
train_ds_names:
- fw_qa_v2_2k_len_level_1_small
- pwc_compact
# these provide exact tokens needed, no need to use self-gen data
- squad_compact
- ropes_compact
- drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

53
configs/small_exp/qa_short_ctx_self_gen_lv0_closed_qa_1_small.yaml
View file

@ -1,53 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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_1.0/fw_qa_v2/min_0_to_2000/train/*level_0.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
# these provide exact tokens needed, no need to use self-gen data
- squad_compact
- ropes_compact
- drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

53
configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_0_small.yaml
View file

@ -1,53 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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/fw_qa_v2/min_0_to_2000/train/*level_1.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
# these provide exact tokens needed, no need to use self-gen data
- squad_compact
- ropes_compact
- drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

38
configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1.yaml Normal file
View file

@ -0,0 +1,38 @@
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
use_kl_loss: true
n_latent_queries: 208
num_blocks: 9
num_self_attn_per_block: 0
gradient_accumulation_steps: 8
max_packed_inp_len: 8192
max_packed_ctx_len: 8192
# data
train_ds_names:
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes
- hotpot_qa
- negative_nq
- booksum
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet

39
configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_and_lv3.yaml Normal file
View file

@ -0,0 +1,39 @@
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
use_kl_loss: true
n_latent_queries: 208
num_blocks: 9
num_self_attn_per_block: 0
gradient_accumulation_steps: 16
max_packed_inp_len: 4096
max_packed_ctx_len: 4096
# data
train_ds_names:
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_3*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes
- hotpot_qa
- negative_nq
- booksum
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet

41
configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_l2l.yaml Normal file
View file

@ -0,0 +1,41 @@
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
use_kl_loss: true
ctx_encoder_type: per_layer_activations
n_latent_queries: 8
num_blocks: 9
num_self_attn_per_block: 0
gradient_accumulation_steps: 16
max_packed_inp_len: 4096
max_packed_ctx_len: 4096
# data
train_ds_names:
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes
- hotpot_qa
- negative_nq
- booksum
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet

37
configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_next_token_pred.yaml Normal file
View file

@ -0,0 +1,37 @@
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
use_kl_loss: false
n_latent_queries: 208
num_blocks: 9
num_self_attn_per_block: 0
gradient_accumulation_steps: 8
max_packed_inp_len: 8192
max_packed_ctx_len: 8192
# data
train_ds_names:
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes
- hotpot_qa
- negative_nq
- booksum
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet

53
configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_small.yaml
View file

@ -1,53 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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_1.0/fw_qa_v2/min_0_to_2000/train/*level_1.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
# these provide exact tokens needed, no need to use self-gen data
- squad_compact
- ropes_compact
- drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

53
configs/small_exp/qa_short_ctx_self_gen_lv3_closed_qa_0.5_small.yaml
View file

@ -1,53 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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.5/fw_qa_v2/min_0_to_2000/train/*level_3.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
# these provide exact tokens needed, no need to use self-gen data
- squad_compact
- ropes_compact
- drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

53
configs/small_exp/qa_short_ctx_self_gen_lv3_closed_qa_1_small.yaml
View file

@ -1,53 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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_1.0/fw_qa_v2/min_0_to_2000/train/*level_3.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
# these provide exact tokens needed, no need to use self-gen data
- squad_compact
- ropes_compact
- drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

49
configs/tiny_exp/fw_qa_v2_level_0_tiny.yaml
View file

@ -1,49 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
# data
train_ds_names:
- fw_qa_v2_2k_len_level_0_tiny
val_ds_names:
- fw_qa_v2_2k_len_level_0
- squad
- drop
- ropes
- pwc

49
configs/tiny_exp/fw_qa_v2_level_1_tiny.yaml
View file

@ -1,49 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
# data
train_ds_names:
- fw_qa_v2_2k_len_level_1_tiny
val_ds_names:
- fw_qa_v2_2k_len_level_0
- squad
- drop
- ropes
- pwc

49
configs/tiny_exp/fw_qa_v2_level_2_tiny.yaml
View file

@ -1,49 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
# data
train_ds_names:
- fw_qa_v2_2k_len_level_2_tiny
val_ds_names:
- fw_qa_v2_2k_len_level_0
- squad
- drop
- ropes
- pwc

49
configs/tiny_exp/fw_qa_v2_level_3_tiny.yaml
View file

@ -1,49 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
# data
train_ds_names:
- fw_qa_v2_2k_len_level_3_tiny
val_ds_names:
- fw_qa_v2_2k_len_level_0
- squad
- drop
- ropes
- pwc

58
configs/tiny_exp/qa_short_ctx_no_fw_qa.yaml
View file

@ -1,58 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
# short ctx data
train_ds_names:
- squad_compact
- pwc_compact
- drop_compact
- ropes_compact
# the contexts are kinda long-ish
# - narrativeqa
# - quoref
# - synthetic_convqa
val_ds_names:
- fw_qa_v2_2k_len_level_0
- squad
- drop
- ropes
- pwc

52
configs/tiny_exp/qa_short_ctx_self_gen_lv0_tiny.yaml
View file

@ -1,52 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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/fw_qa_v2_2k_len_level_0_tiny
- self_gen/google/gemma-2-2b-it_temp_0.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

53
configs/tiny_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_tiny.yaml
View file

@ -1,53 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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_1.0/fw_qa_v2_2k_len_level_1_tiny
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
# these provide exact tokens needed, no need to use self-gen data
- squad_compact
- ropes_compact
- drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

52
configs/tiny_exp/qa_short_ctx_self_gen_lv1_tiny.yaml
View file

@ -1,52 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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/fw_qa_v2_2k_len_level_1_tiny
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
- squad_compact
- ropes_compact
- drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

52
configs/tiny_exp/qa_short_ctx_self_gen_lv1_tiny_temp_0.3.yaml
View file

@ -1,52 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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.3/fw_qa_v2_2k_len_level_1_tiny
- self_gen/google/gemma-2-2b-it_temp_0.3/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.3/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.3/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.3/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

52
configs/tiny_exp/qa_short_ctx_self_gen_lv1_tiny_temp_0.5.yaml
View file

@ -1,52 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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.5/fw_qa_v2_2k_len_level_1_tiny
- self_gen/google/gemma-2-2b-it_temp_0.5/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.5/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.5/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.5/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

52
configs/tiny_exp/qa_short_ctx_self_gen_lv1_tiny_temp_0.7.yaml
View file

@ -1,52 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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.7/fw_qa_v2_2k_len_level_1_tiny
- self_gen/google/gemma-2-2b-it_temp_0.7/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.7/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.7/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.7/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

52
configs/tiny_exp/qa_short_ctx_self_gen_lv1_tiny_temp_1.0.yaml
View file

@ -1,52 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
# data
train_ds_names:
- self_gen/google/gemma-2-2b-it_temp_1.0/fw_qa_v2_2k_len_level_1_tiny
- self_gen/google/gemma-2-2b-it_temp_1.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_1.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_1.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

52
configs/tiny_exp/qa_short_ctx_self_gen_lv2_tiny.yaml
View file

@ -1,52 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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/fw_qa_v2_2k_len_level_2_tiny
- self_gen/google/gemma-2-2b-it_temp_0.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

52
configs/tiny_exp/qa_short_ctx_self_gen_lv3_closed_qa_1_tiny.yaml
View file

@ -1,52 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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_1.0/fw_qa_v2_2k_len_level_3_tiny
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
- squad_compact
- ropes_compact
- drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

52
configs/tiny_exp/qa_short_ctx_self_gen_lv3_tiny.yaml
View file

@ -1,52 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# 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/fw_qa_v2/min_0_to_2000/train/ds_013*level_3*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

59
configs/tiny_exp/qa_short_ctx_tiny.yaml
View file

@ -1,59 +0,0 @@
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: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 16
dataloader_num_workers: 8
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
# short ctx data
train_ds_names:
- fw_qa_v2_2k_len_level_3_tiny
- squad_compact
- pwc_compact
- drop_compact
- ropes_compact
# the contexts are kinda long-ish
# - narrativeqa
# - quoref
# - synthetic_convqa
val_ds_names:
- fw_qa_v2_2k_len_level_0
- squad
- drop
- ropes
- pwc

20
data/download_facts_grounding.py Normal file
View file

@ -0,0 +1,20 @@
# Install dependencies as needed:
# pip install kagglehub[hf-datasets]
import kagglehub
from kagglehub import KaggleDatasetAdapter
# Set the path to the file you'd like to load
file_path = "data/raw_datasets/FACTS/ds.parquet"
# Load the latest version
hf_dataset = kagglehub.load_dataset(
KaggleDatasetAdapter.HUGGING_FACE,
"deepmind/FACTS-grounding-examples",
file_path,
# Provide any additional arguments like
# sql_query, hf_kwargs, or pandas_kwargs. See
# the documenation for more information:
# https://github.com/Kaggle/kagglehub/blob/main/README.md#kaggledatasetadapterhugging_face
)
print("Hugging Face Dataset:", hf_dataset)

220535
data/eval/facts/examples.csv Normal file
View file

File diff suppressed because one or more lines are too long

2
data/generate_negative_natural_questions.py
View file

@ -6,7 +6,7 @@ from tqdm import tqdm
if __name__ == "__main__":
ds = load_dataset("google-research-datasets/natural_questions")
ds = ds.shuffle(seed=42)
for split in ds:
for split in ["validation"]:
out = []
for i, sample in enumerate(tqdm(ds[split])):
ctx = " ".join(

56
data/self_generate_qa.py
View file

@ -91,6 +91,29 @@ MODEL_CTX_LEN = {
}
def truncate_middle_if_too_long(
input_ids: list[int],
max_length: int,
max_new_tokens: int = 256,
) -> list[int]:
"""
Truncate the middle of a list of tokens to fit within a maximum length.
Args:
tokens: List of token IDs
max_length: Maximum length for the truncated tokens
Returns:
List of truncated token IDs
"""
max_new_tokens_half = max_new_tokens // 2
# leave max_new_tokens for generation
half = max_length // 2 - max_new_tokens_half
if len(input_ids) > max_length:
return input_ids[:half] + input_ids[-half:]
return input_ids
def get_prompt(context: str, q: str) -> str:
return PROMPT_TEMPLATE.format(context=context, question=q)
@ -197,6 +220,7 @@ def self_generate(
llm: LLM,
system_template: str,
parquet_file: str | None = None,
do_truncate: bool = False,
) -> None:
"""Process a single dataset and generate QA pairs."""
@ -294,6 +318,28 @@ def self_generate(
chunk_messages = create_messages(
chunk_ctxs, chunk_questions, args.vllm_model, SELF_GEN_SYSTEM_MSG
)
if do_truncate:
# we should only do this for evaluation data
tokenized_contents = tk(
[m[0]["content"] for m in chunk_messages],
add_special_tokens=False,
return_attention_mask=False,
)
tokenized_contents["input_ids"] = [
truncate_middle_if_too_long(
ids,
max_length=MODEL_CTX_LEN[args.vllm_model],
max_new_tokens=1024,
)
for ids in tokenized_contents["input_ids"]
]
contents = tk.batch_decode(
tokenized_contents["input_ids"], skip_special_tokens=True
)
for c, m in zip(contents, chunk_messages):
m[0]["content"] = c
print(f"Generating from {len(chunk_messages)} contexts")
# Clear GPU memory before processing the next chunk
@ -309,7 +355,7 @@ def self_generate(
sys_tokens,
n_sys_tokens,
chunk_ctxs,
ds["ctx_ids"][start : start + chunk_size],
ds[start : start + chunk_size]["ctx_ids"],
chunk_questions,
chunk_messages,
k,
@ -560,6 +606,11 @@ def parse_args() -> argparse.Namespace:
default=0.0,
help="Probability of using closed QA prompt template (default: 0.0)",
)
parser.add_argument(
"--do_truncate",
action="store_true",
help="Truncate contexts to fit model context length",
)
return parser.parse_args()
@ -599,7 +650,7 @@ 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)
self_generate(ds_name, split, args, llm, SYSTEM_TEMPLATE, args.do_truncate)
else:
assert args.glob_pattern, (
"glob_pattern must be provided if no ds_names or config"
@ -614,4 +665,5 @@ if __name__ == "__main__":
args=args,
llm=llm,
system_template=SYSTEM_TEMPLATE,
do_truncate=args.do_truncate,
)

2
pyproject.toml
View file

@ -32,6 +32,8 @@ dependencies = [
"vllm==0.8.5",
"huggingface-hub[hf-transfer]>=0.32.0",
"opt-einsum>=3.4.0",
"kagglehub[hf-datasets]>=0.3.12",
"kaggle>=1.7.4.5",
]
[build-system]

6
run_eval.py
View file

@ -51,6 +51,12 @@ if __name__ == "__main__":
default=32,
help="Eval batch size for generation",
)
parser.add_argument(
"--max_new_tokens",
type=int,
default=256,
help="Maximum number of new tokens to generate during evaluation",
)
parser.add_argument(
"--remove_context",
action="store_true",

34
scripts/short_ctx/gemma_fw_qa_level_0.sh
View file

@ -1,34 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29562 intx_sft.py configs/fw_qa_v2_level_0.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=5 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--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 \
--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=50

34
scripts/short_ctx/gemma_fw_qa_level_0_tiny.sh
View file

@ -1,34 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29562 intx_sft.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=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--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 \
--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=50

34
scripts/short_ctx/gemma_fw_qa_level_1.sh
View file

@ -1,34 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29563 intx_sft.py configs/fw_qa_v2_level_1.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=5 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--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 \
--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=50

34
scripts/short_ctx/gemma_fw_qa_level_2.sh
View file

@ -1,34 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29564 intx_sft.py configs/fw_qa_v2_level_2.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=5 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--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 \
--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=50

34
scripts/short_ctx/gemma_fw_qa_level_3.sh
View file

@ -1,34 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29565 intx_sft.py configs/fw_qa_v2_level_3.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=5 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--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 \
--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=50

34
scripts/short_ctx/gemma_fw_qa_level_3_tiny.sh
View file

@ -1,34 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29563 intx_sft.py configs/fw_qa_v2_level_3_tiny.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=5 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--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 \
--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=50

34
scripts/short_ctx/gemma_qa_short_ctx.sh
View file

@ -1,34 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29571 intx_sft.py configs/qa_short_ctx.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=2 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--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 \
--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=50

46
scripts/short_ctx/gemma_qa_short_ctx_exp_default.sh
View file

@ -9,33 +9,21 @@
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
# Default arguments - these can be overridden by passing arguments to this script
default_args=(
"--model_name_or_path=google/gemma-2-2b-it"
"--max_steps=50_000"
"--target_modules=down_proj"
"--lora_r=8"
"--eval_strategy=no"
"--max_qas_len=2048"
"--max_qas_per_sample=1"
"--per_rank_gen=True"
"--per_layer_processing=True"
"--gen_lora_l1_reg_coef=0.1"
)
# Pass all script arguments to the training command
# $1 comes first, then defaults, then remaining arguments (which can override defaults)
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=12_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=5 \
--per_device_eval_batch_size=64 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=0 \
--n_latent_queries=208 \
--num_pre_head_layers=1 \
--lora_r=8 \
--eval_strategy=no \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=2048 \
--max_packed_inp_len=8192 \
--max_packed_ctx_len=16384 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50
--num_processes=4 --gpu_ids all train.py $1 "${default_args[@]}" "${@:2}"

42
scripts/short_ctx/gemma_qa_short_ctx_exp_default_max_qas_per_sample_1.sh
View file

@ -1,42 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=12_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=5 \
--per_device_eval_batch_size=64 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=0 \
--n_latent_queries=208 \
--num_pre_head_layers=1 \
--lora_r=8 \
--eval_strategy=no \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_per_sample=1 \
--max_qas_len=2048 \
--max_packed_inp_len=8192 \
--max_packed_ctx_len=16384 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50

40
scripts/short_ctx/gemma_qa_short_ctx_exp_default_more_prehead.sh
View file

@ -1,40 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=12_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=5 \
--per_device_eval_batch_size=64 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=0 \
--n_latent_queries=208 \
--num_pre_head_layers=4 \
--lora_r=8 \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=2048 \
--max_packed_inp_len=8192 \
--max_packed_ctx_len=16384 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50

41
scripts/short_ctx/gemma_qa_short_ctx_exp_default_more_prehead_max_qas_per_sample_1.sh
View file

@ -1,41 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=12_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=5 \
--per_device_eval_batch_size=64 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=0 \
--n_latent_queries=208 \
--num_pre_head_layers=4 \
--lora_r=8 \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_per_sample=1 \
--max_qas_len=2048 \
--max_packed_inp_len=8192 \
--max_packed_ctx_len=16384 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50

40
scripts/short_ctx/gemma_qa_short_ctx_exp_self_attn.sh
View file

@ -1,40 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=12_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--target_modules=down_proj \
--num_self_attends_per_block=8 \
--n_cross_attn_layers=1 \
--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 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=2048 \
--max_packed_inp_len=8192 \
--max_packed_ctx_len=16384 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50

34
scripts/short_ctx/gemma_qa_short_ctx_no_fw_qa.sh
View file

@ -1,34 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29581 intx_sft.py configs/qa_short_ctx_no_fw_qa.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=2 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--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=1e-4 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--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=50

40
scripts/short_ctx/gemma_qa_short_ctx_old_conf.sh
View file

@ -1,40 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((29560 + ($SLURM_JOBID % 10)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=3 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=8 \
--per_device_eval_batch_size=64 \
--target_modules=down_proj \
--num_blocks=9 \
--num_self_attn_per_block=0 \
--n_latent_queries=208 \
--num_pre_head_layers=1 \
--lora_r=8 \
--eval_strategy=no \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--use_sequence_packing=True \
--max_qas_len=2048 \
--max_qas_per_sample=1 \
--max_packed_inp_len=8192 \
--max_packed_ctx_len=8192 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50

42
scripts/short_ctx/gemma_qa_short_ctx_old_conf_8_gpus.sh
View file

@ -1,42 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=8
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=5 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=4 \
--per_device_eval_batch_size=64 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=0 \
--n_latent_queries=208 \
--num_pre_head_layers=4 \
--lora_r=8 \
--eval_strategy=no \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=2048 \
--max_qas_per_sample=1 \
--max_packed_inp_len=8192 \
--max_packed_ctx_len=16384 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50

41
scripts/short_ctx/gemma_qa_short_ctx_old_conf_distill.sh
View file

@ -1,41 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((29560 + ($SLURM_JOBID % 10)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=3 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=16 \
--per_device_eval_batch_size=64 \
--target_modules=down_proj \
--num_blocks=9 \
--num_self_attn_per_block=0 \
--n_latent_queries=208 \
--num_pre_head_layers=1 \
--lora_r=8 \
--eval_strategy=no \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--use_sequence_packing=True \
--max_qas_len=2048 \
--max_qas_per_sample=1 \
--max_packed_inp_len=4096 \
--max_packed_ctx_len=4096 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50 \
--use_kl_loss=True

42
scripts/short_ctx/gemma_qa_short_ctx_old_conf_l2l_distill.sh
View file

@ -1,42 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((29560 + ($SLURM_JOBID % 10)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=3 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=16 \
--per_device_eval_batch_size=64 \
--target_modules=down_proj \
--num_blocks=9 \
--num_self_attn_per_block=0 \
--ctx_encoder_type=per_layer_activations \
--n_latent_queries=8 \
--num_pre_head_layers=1 \
--lora_r=8 \
--eval_strategy=no \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--use_sequence_packing=True \
--max_qas_len=2048 \
--max_qas_per_sample=1 \
--max_packed_inp_len=4096 \
--max_packed_ctx_len=4096 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50 \
--use_kl_loss=True

41
scripts/short_ctx/gemma_qa_short_ctx_old_conf_more_grad_acc.sh
View file

@ -1,41 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=12_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=16 \
--per_device_eval_batch_size=64 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=0 \
--n_latent_queries=208 \
--num_pre_head_layers=4 \
--lora_r=8 \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=2048 \
--max_qas_per_sample=1 \
--max_packed_inp_len=8192 \
--max_packed_ctx_len=16384 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50

41
scripts/short_ctx/gemma_qa_short_ctx_perceiver.sh
View file

@ -1,41 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=30_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=4 \
--per_device_eval_batch_size=64 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=3 \
--n_latent_queries=208 \
--num_pre_head_layers=1 \
--lora_r=8 \
--eval_strategy=no \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=2048 \
--max_packed_inp_len=8192 \
--max_packed_ctx_len=16384 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50

41
scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l.sh
View file

@ -1,41 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=30_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=8 \
--per_device_eval_batch_size=32 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=3 \
--ctx_encoder_type=per_layer_activations \
--n_latent_queries=8 \
--num_pre_head_layers=1 \
--lora_r=8 \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=2048 \
--max_packed_inp_len=4096 \
--max_packed_ctx_len=8192 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.0 \
--logging_steps=50

41
scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l_128_max_qas_len.sh
View file

@ -1,41 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=50_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=4 \
--per_device_eval_batch_size=32 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=3 \
--ctx_encoder_type=per_layer_activations \
--n_latent_queries=8 \
--num_pre_head_layers=1 \
--lora_r=8 \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=128 \
--max_packed_inp_len=4096 \
--max_packed_ctx_len=4096 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.0 \
--logging_steps=50

41
scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l_32_max_qas_len.sh
View file

@ -1,41 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=50_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=4 \
--per_device_eval_batch_size=32 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=3 \
--ctx_encoder_type=per_layer_activations \
--n_latent_queries=8 \
--num_pre_head_layers=1 \
--lora_r=8 \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=32 \
--max_packed_inp_len=4096 \
--max_packed_ctx_len=4096 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.0 \
--logging_steps=50

41
scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l_512_max_qas_len.sh
View file

@ -1,41 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=50_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=4 \
--per_device_eval_batch_size=32 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=3 \
--ctx_encoder_type=per_layer_activations \
--n_latent_queries=8 \
--num_pre_head_layers=1 \
--lora_r=8 \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=512 \
--max_packed_inp_len=4096 \
--max_packed_ctx_len=4096 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50

41
scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l_64_max_qas_len.sh
View file

@ -1,41 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=50_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=4 \
--per_device_eval_batch_size=32 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=3 \
--ctx_encoder_type=per_layer_activations \
--n_latent_queries=8 \
--num_pre_head_layers=1 \
--lora_r=8 \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=64 \
--max_packed_inp_len=4096 \
--max_packed_ctx_len=4096 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.0 \
--logging_steps=50

42
scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l_max_qas_per_sample_1.sh
View file

@ -1,42 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=30_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=8 \
--per_device_eval_batch_size=32 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=3 \
--ctx_encoder_type=per_layer_activations \
--n_latent_queries=8 \
--num_pre_head_layers=1 \
--lora_r=8 \
--eval_steps=1000 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=2048 \
--max_qas_per_sample=1 \
--max_packed_inp_len=4096 \
--max_packed_ctx_len=8192 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.0 \
--logging_steps=50

41
scripts/short_ctx/gemma_qa_short_ctx_perceiver_max_qas_per_sample_1.sh
View file

@ -1,41 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
uv run accelerate launch --main_process_port $port \
--num_processes=4 --gpu_ids all intx_sft.py $1 \
--model_name_or_path=google/gemma-2-2b-it \
--max_steps=10_000 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--target_modules=down_proj \
--num_blocks=8 \
--num_self_attn_per_block=3 \
--n_latent_queries=208 \
--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 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_qas_len=2048 \
--max_qas_per_sample=1 \
--max_packed_inp_len=8192 \
--max_packed_ctx_len=16384 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50

34
scripts/short_ctx/gemma_qa_short_ctx_self_gen.sh
View file

@ -1,34 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=8
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
uv run accelerate launch --num_processes=8 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29572 intx_sft.py configs/qa_short_ctx_self_gen.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=2 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=4 \
--per_device_eval_batch_size=64 \
--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 \
--add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True \
--max_packed_inp_len=4096 \
--max_packed_ctx_len=8192 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--logging_steps=50

34
scripts/short_ctx/gemma_qa_short_ctx_self_gen_no_fw_qa.sh
View file

@ -1,34 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29581 intx_sft.py configs/qa_short_ctx_self_gen_no_fw_qa.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=2 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--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=1e-4 \
--lora_dropout=0.0 \
--neftune_noise_alpha=5 \
--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=50

34
scripts/short_ctx/gemma_squad.sh
View file

@ -1,34 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29561 intx_sft.py configs/squad.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=5 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--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 \
--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=50

34
scripts/short_ctx/gemma_squad_and_fw_qa_level_3.sh
View file

@ -1,34 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29566 intx_sft.py configs/squad_fw_qa_v2_level_3.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=5 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=2 \
--per_device_eval_batch_size=64 \
--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 \
--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=50

62
src/ctx_to_lora/configs.py
View file

@ -133,10 +133,25 @@ class ExperimentSetup(str, Enum):
@dataclass
class TrainingArguments(TrainingArguments):
output_dir: str = field(
default="",
metadata={"help": "Placeholder. Will be overwritten by train.py"},
)
tf32: bool = field(
default=True,
metadata={"help": "Whether to use tf32 precision."},
)
bf16: bool = field(
default=True,
metadata={"help": "Whether to use bf16 precision."},
)
label_names: list[str] = field(
default=("labels",),
metadata={
"help": "List of strings to specify the label names in the dataset. "
"This is used to compute the loss and metrics."
},
)
include_for_metrics: list[str] = field(
default=("inputs",),
metadata={
@ -144,6 +159,21 @@ class TrainingArguments(TrainingArguments):
"Options: 'inputs', 'loss'."
},
)
per_device_eval_batch_size: int = field(
default=64,
metadata={
"help": "Batch size for evaluation. "
"If not set, will use the same as per_device_train_batch_size."
},
)
per_device_train_batch_size: int = field(
default=1,
metadata={
"help": "Batch size for training. "
"If not set, will use the same as per_device_eval_batch_size."
},
)
# TODO: use this! (check trainer.py for proper computation)
average_tokens_across_devices: bool = field(
default=False,
metadata={"help": "compute num_items_in_batch across devices."},
@ -165,6 +195,18 @@ class TrainingArguments(TrainingArguments):
default=8,
metadata={"help": "Number of subprocesses to use for data loading."},
)
neftune_noise_alpha: float = field(
default=5.0,
metadata={"help": "Neftune noise alpha for the optimizer."},
)
learning_rate: float = field(
default=4e-5,
metadata={"help": "Initial learning rate."},
)
weight_decay: float = field(
default=0.01,
metadata={"help": "Weight decay for the optimizer."},
)
optim: str = field(
default="adamw_torch_fused",
metadata={"help": "Optimizer."},
@ -202,7 +244,7 @@ class TrainingArguments(TrainingArguments):
metadata={"help": "Evaluation strategy."},
)
eval_steps: int = field(
default=10_000,
default=1_000,
metadata={"help": "Evaluation steps."},
)
metric_for_best_model: str = field(
@ -225,7 +267,7 @@ class TrainingArguments(TrainingArguments):
default="steps",
)
save_steps: int = field(
default=10_000,
default=5_000,
)
save_safetensors: bool = field(
default=False,
@ -316,9 +358,13 @@ class CtxTrainingArguments:
metadata={"help": "Maximum context length for training."},
)
use_sequence_packing: bool = field(
default=False,
default=True,
metadata={"help": "Whether to use sequence packing."},
)
max_ctx_len: int = field(
default=-1,
metadata={"help": "Max context length. Overrides ctx tokenizer length."},
)
max_qas_len: int = field(
default=2**11,
metadata={
@ -344,7 +390,7 @@ class CtxTrainingArguments:
)
max_new_tokens: int | None = field(
default=2**10,
default=256,
metadata={"help": "Maximum new tokens for generation-based evaluation."},
)
gen_per_device_eval_batch_size: int | None = field(
@ -359,6 +405,10 @@ class CtxTrainingArguments:
default=False,
metadata={"help": "Whether to use KL loss."},
)
use_per_ctx_average_loss: bool = field(
default=False,
metadata={"help": "Whether to use per-context average loss."},
)
gen_lora_l1_reg_coef: float = field(
default=0.0,
metadata={"help": "L1 regularization coefficient for generated LoRAs."},
@ -429,7 +479,7 @@ class HypernetArguments:
metadata={"help": "Whether to use token mixing block."},
)
num_pre_head_layers: int = field(
default=4, metadata={"help": "# of layers before hypernet head"}
default=1, metadata={"help": "# of layers before hypernet head"}
)
@ -499,7 +549,7 @@ class AggregatorArguments:
metadata={"help": "Number of blocks for Perceiver."},
)
num_self_attn_per_block: int = field(
default=6,
default=0,
metadata={"help": "Number of self-attention layers per block for Perceiver."},
)
shared_weights: bool = field(

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

@ -546,12 +546,29 @@ DS_KWARGS = {
),
# for negative training and evaluating
"negative_nq": dict(
validation=dict(
path="json",
data_files="data/raw_datasets/negative_natural_questions/validation.jsonl",
split="train[:1000]",
),
test=dict(
path="json",
data_files="data/raw_datasets/negative_natural_questions/validation.jsonl",
split="train",
),
),
"facts": dict(
validation=dict(
path="csv",
data_files="data/eval/facts/examples.csv",
split="train[:400]",
),
test=dict(
path="csv",
data_files="data/eval/facts/examples.csv",
split="train",
),
),
}
# LongBench kwargs
@ -630,12 +647,22 @@ CLOSED_QA_INTX_TEMPLATES = [
EVAL_INTX_TEMPLATES = {
# binary (yes/no, a/b) qa given ctx
"ropes": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
# short-ctx reasoning
"drop": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
"negative_nq": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
"triviaqa_retrieved": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
"hotpot_qa": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
# short-ctx extractive qa
"squad": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
# retrieved noisy ctx?
"triviaqa_retrieved": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
# doc w/ distractors qa
"hotpot_qa": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
# unrelated ctx-qa
"negative_nq": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
# summary
"booksum": "Write a summary for the story. Make it a *one-paragraph monologue*. Don't use bullet point. Do not output anything else beyond the summary.",
"gov_report": "Write a summary for the report. Make it a *one-paragraph monologue*. Don't use bullet point. Do not output anything else beyond the summary.",
# longbench
"longbench/qasper": 'Answer the question as concisely as you can, using a single phrase or sentence if possible.\nIf the question cannot be answered based on the information in the article, write "unanswerable".\nIf the question is a yes/no question, answer "yes", "no", or "unanswerable". Do not provide any explanation.\n\nQuestion: {input}',
"longbench/narrativeqa": "Answer the question as concisely as you can, using a single phrase if possible. Do not provide any explanation.\n\nQuestion: {input}",
"longbench/multifieldqa_en": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",

19
src/ctx_to_lora/data/preprocessing_fn.py
View file

@ -55,6 +55,25 @@ def get_preprocessing_fn(
"responses": responses[:min_len],
}
elif ds_name.startswith("ctx_qa"):
def f(sample):
ctx = sample["context"]
# remove unintended prompt
ctx = ctx.split("\n...\n\n### Code ###\n")[-1]
sample["context"] = ctx
return sample
elif ds_name == "facts":
def f(sample):
prompt = sample["system_instruction"] + sample["user_request"]
return {
"context": sample["context_document"],
"prompt": prompt,
"response": "PLACEHOLDER. PLEASE GENERATE SELF-GEN DATA.",
}
elif ds_name.startswith("longbench"):
def f(sample):

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

@ -170,12 +170,12 @@ def load_and_process_dataset(
ds = ds.map(
get_preprocessing_fn(ds_name, is_eval),
remove_columns=cols_to_remove,
num_proc=num_proc,
num_proc=16,
)
ds = ds.filter(
filter_none,
batched=False,
num_proc=num_proc,
num_proc=16,
)
return ds
@ -191,6 +191,7 @@ def get_tokenized_dataset(
ctx_tokenizer: PreTrainedTokenizerBase,
add_ctx_to_chat: bool,
use_kl_loss: bool,
max_new_tokens: int = 256,
set_format: str | None = None,
) -> dict[str, Any]:
if max_qas_len > 0:
@ -210,9 +211,10 @@ def get_tokenized_dataset(
base_model_max_len=base_model_max_len,
ctx_model_max_len=ctx_model_max_len,
add_ctx_to_chat=add_ctx_to_chat,
use_kl_loss=use_kl_loss,
use_kl_loss=True, # HACK: for consistent hash (to be removed)
need_ctx_ids=need_ctx_ids,
split=split,
max_new_tokens=max_new_tokens,
set_format=set_format,
)
@ -229,6 +231,10 @@ def get_tokenized_dataset(
# load the cached ds
logger.info(f"Loaded tokenized dataset from {ds_path}")
tokenized_ds = datasets.load_from_disk(ds_path)
if (not use_kl_loss) and ("logprobs_vals" in tokenized_ds.column_names):
tokenized_ds = tokenized_ds.remove_columns(
["logprobs_vals", "logprobs_indices"]
)
return tokenized_ds
num_proc = 4
@ -236,7 +242,17 @@ def get_tokenized_dataset(
**load_and_process_kwargs,
num_proc=num_proc,
)
if use_kl_loss:
if "train" in split and "logprobs_vals" not in ds.column_names:
raise ValueError(
"`use_kl_loss` is set to True but 'logprobs_vals' column "
"is not present in the dataset."
)
logger.info(f"Constructing and tokenizing {ds_name} with {split} split...")
# HACK
tokenize_kwargs.pop("use_kl_loss", None)
tokenized_ds = construct_and_tokenize_ctx_qa(
ds=ds,
tokenizer=tokenizer,
@ -247,9 +263,14 @@ def get_tokenized_dataset(
if ("train" in split) and is_caching_enabled():
tokenized_ds = tokenized_ds.shuffle()
tokenized_ds.save_to_disk(ds_path, num_proc=num_proc)
tokenized_ds.save_to_disk(ds_path, num_proc=16)
# force reload from disk for fingerprint consistency
tokenized_ds = datasets.load_from_disk(ds_path)
if (not use_kl_loss) and ("logprobs_vals" in tokenized_ds.column_names):
tokenized_ds = tokenized_ds.remove_columns(
["logprobs_vals", "logprobs_indices"]
)
return tokenized_ds
@ -261,34 +282,24 @@ def construct_and_tokenize_ctx_qa(
ctx_model_max_len,
ctx_tokenizer,
add_ctx_to_chat,
use_kl_loss,
need_ctx_ids,
ds,
split,
max_new_tokens,
set_format=None,
num_proc=None,
):
if use_kl_loss:
if "train" in split and "logprobs_vals" not in ds.column_names:
raise ValueError(
"`use_kl_loss` is set to True but 'logprobs_vals' column "
"is not present in the dataset."
)
else:
if "logprobs_vals" in ds.column_names:
ds = ds.remove_columns(["logprobs_vals", "logprobs_indices"])
# for sft + chat_model, we need to convert the dataset to chat format
if "input_ids" in ds.column_names and "response_start_end" in ds.column_names:
# already tokenized dataset (e.g., self-gen qa data)
tokenized_ds = ds.map(get_labels_from_input_ids, num_proc=num_proc)
tokenized_ds = ds.map(get_labels_from_input_ids, num_proc=16)
else:
# construct messages from prompts and responses
# add "messages_list" field
ds = ds.map(
convert_ctx_prompt_response_to_messages,
fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat},
num_proc=num_proc,
num_proc=16,
)
# add `input_ids`, `attention_mask`, `labels`
os.environ["TOKENIZERS_PARALLELISM"] = "true"
@ -302,6 +313,10 @@ def construct_and_tokenize_ctx_qa(
tokenized_ds = tokenized_ds.remove_columns(
[col for col in tokenized_ds.column_names if col not in COLS_TO_KEEP_TOKENIZED],
)
tokenized_ds = tokenized_ds.filter(
lambda x: bool(x["input_ids"]), # remove empty "input_ids"
num_proc=num_proc,
)
if need_ctx_ids:
if "ctx_ids" not in tokenized_ds.column_names:
@ -321,7 +336,7 @@ def construct_and_tokenize_ctx_qa(
tokenized_ds = tokenized_ds.filter(
len_filter,
fn_kwargs={"max_length": ctx_model_max_len, "keys": ["ctx_ids"]},
num_proc=num_proc,
num_proc=16,
)
logging.info(f"Split too long QAs with max length {max_qas_len}")
@ -332,8 +347,8 @@ def construct_and_tokenize_ctx_qa(
"max_qas_per_sample": max_qas_per_sample,
},
batched=True,
batch_size=100_000,
num_proc=4,
batch_size=12_500,
num_proc=16,
)
if "train" not in split:
# squeeze since we always have one query per sample in eval
@ -343,7 +358,7 @@ def construct_and_tokenize_ctx_qa(
fn_kwargs={
"max_length": base_model_max_len,
"columns": ["input_ids", "labels"],
"max_new_tokens": 256,
"max_new_tokens": max_new_tokens,
},
)
tokenized_ds = tokenized_ds.map(
@ -589,7 +604,7 @@ def split_too_long_qas(
target_vals, target_indices = None, None
inp_len = len(inp_ids)
if inp_len > max_qas_len:
if (max_qas_len > 0) and (inp_len > max_qas_len):
# Skip individual sequences that are too long
n_skip += 1
continue
@ -774,7 +789,7 @@ def pack(
"metadata_path": f"{ds_path}/packing_metadata.json",
},
batched=True,
batch_size=250_000,
batch_size=125_000,
num_proc=num_proc,
remove_columns=ds.column_names,
)

95
src/ctx_to_lora/eval_utils.py
View file

@ -20,6 +20,7 @@ from transformers import (
Trainer,
set_seed,
)
from transformers.trainer_utils import find_executable_batch_size
from ctx_to_lora.data.collator import eval_collator, generation_collator
from ctx_to_lora.data.definitions import (
@ -50,13 +51,76 @@ logger = logging.getLogger()
sys.modules["ctx_to_lora.modeling_utils"] = hypernet
# from https://github.com/huggingface/transformers/issues/32215#issuecomment-2386274138
class EvaluateAutoFindBatchSizeMixin:
def evaluate(self, *args, **kwargs):
print(f"Auto-finding batch size for evaluate")
# can't set self.args.eval_batch_size, it is computed from self.args.per_device_eval_batch_size
super_evaluate = super().evaluate # super cannot be called inside lambda
# assert not self.args.dataloader_persistent_workers, "dataloader_persistent_workers not supported because eval batch size does not change"
def new_evaluate(batch_size):
old_eval_batch_size = self.args.eval_batch_size
self.args.per_device_eval_batch_size = batch_size
if self.args.eval_batch_size != old_eval_batch_size:
if self.args.dataloader_persistent_workers and hasattr(
self, "_eval_dataloaders"
):
del self._eval_dataloaders # to shutdown the workers
logger.info(
f"Setting eval_batch_size to {self.args.eval_batch_size} (per device: {batch_size}) from old {old_eval_batch_size}"
)
return super_evaluate(*args, **kwargs)
# don't reset at end, but keep batch size for the next time evaluate is called
return find_executable_batch_size(
new_evaluate,
self.args.per_device_eval_batch_size,
auto_find_batch_size=self.args.auto_find_batch_size,
)()
def predict(self, *args, **kwargs):
print(f"Auto-finding batch size for predict")
# can't set self.args.eval_batch_size, it is computed from self.args.per_device_eval_batch_size
super_predict = super().predict # super cannot be called inside lambda
def new_predict(batch_size):
old_eval_batch_size = self.args.eval_batch_size
self.args.per_device_eval_batch_size = batch_size
if self.args.eval_batch_size != old_eval_batch_size:
if self.args.dataloader_persistent_workers and hasattr(
self, "_eval_dataloaders"
):
del self._eval_dataloaders # to shutdown the workers
logger.info(
f"Setting eval_batch_size to {self.args.eval_batch_size} (per device: {batch_size}) from old {old_eval_batch_size}"
)
return super_predict(*args, **kwargs)
return find_executable_batch_size(
new_predict,
self.args.per_device_eval_batch_size,
auto_find_batch_size=self.args.auto_find_batch_size,
)()
class CustomSeq2SeqTrainer(Seq2SeqTrainer, EvaluateAutoFindBatchSizeMixin):
pass
class CustomTrainer(Trainer, EvaluateAutoFindBatchSizeMixin):
pass
# ============================================================================
# Metrics and Evaluation Utilities
# ============================================================================
# from https://gist.github.com/cloneofsimo/8abd0284d4738f28f04200628f9a83f5
# https://github.com/Nordth/humanize-ai-lib/blob/main/src/humanize-string.ts
import re
_HIDDEN_CHARS = re.compile(
r"[\u00AD\u180E\u200B-\u200F\u202A-\u202E\u2060\u2066-\u2069\uFEFF]"
@ -501,6 +565,7 @@ def eval_generation(
if remove_context:
split_name += "_no_context"
clear_gpu()
eval_result = eval_trainer.predict(
ds,
metric_key_prefix=split_name,
@ -662,6 +727,7 @@ def evaluate(
eval_batch_size: int,
args: Namespace,
split: str,
max_new_tokens: int,
generative: bool,
) -> dict[str, dict]:
"""Main evaluation function."""
@ -691,25 +757,6 @@ def evaluate(
use_flash_attn=True,
)
# if is_liger_kernel_available():
# from liger_kernel.transformers import _apply_liger_kernel_to_instance
# if isinstance(model, ModulatedPretrainedModel):
# print("Applying liger-kernel to ModulatedPretrainedModel")
# if isinstance(model.base_model, PeftModel):
# _apply_liger_kernel_to_instance(model=model.base_model.base_model.model)
# else:
# _apply_liger_kernel_to_instance(model=model.base_model.model)
# if ctx_name is not None:
# print("Applying liger-kernel to ctx_encoder_model")
# _apply_liger_kernel_to_instance(model=model.ctx_encoder.base_model)
# elif isinstance(model, PeftModel):
# print("Applying liger-kernel to PeftModel")
# _apply_liger_kernel_to_instance(model=model.base_model.model)
# elif isinstance(model, PreTrainedModel):
# print("Applying liger-kernel to PretrainedModel")
# _apply_liger_kernel_to_instance(model=model.model)
tokenizer = get_tokenizer(args.model_name_or_path, train=False)
if tokenizer.pad_token_id is None:
tokenizer.pad_token_id = tokenizer.eos_token_id
@ -744,6 +791,7 @@ def evaluate(
ctx_tokenizer=ctx_tokenizer,
add_ctx_to_chat=add_ctx_to_chat,
use_kl_loss=False,
max_new_tokens=max_new_tokens,
set_format="pt",
)
@ -774,7 +822,7 @@ def evaluate(
gen_kwargs = dict(
do_sample=False,
max_new_tokens=256,
max_new_tokens=max_new_tokens,
)
eval_trainer_args = {}
@ -832,12 +880,13 @@ def evaluate(
# spents a few hours on this but couldn't find the reason
for ds_name, ds in datasets.items():
eval_trainer = Trainer(**trainer_kwargs)
clear_gpu()
metrics = eval_teacher_forcing(
eval_trainer, {ds_name: ds}, split, args.remove_context
)
out.update(metrics)
else:
eval_trainer = Seq2SeqTrainer(**trainer_kwargs)
eval_trainer = CustomSeq2SeqTrainer(**trainer_kwargs)
metrics = eval_generation(
eval_trainer,
tokenizer,
@ -861,6 +910,7 @@ def run_eval(
split: str = "validation",
eval_batch_size: int = 8,
remove_context: bool = False,
max_new_tokens: int = 256,
generative: bool = False,
) -> None:
"""Run evaluation with the specified parameters."""
@ -925,5 +975,6 @@ def run_eval(
eval_batch_size,
args,
split,
max_new_tokens,
generative=generative,
)

13
src/ctx_to_lora/modeling/ctx_encoder.py
View file

@ -53,12 +53,13 @@ def maybe_add_batch_dim(kwargs):
class EarlyExit(nn.Module):
def __init__(self, base_model: PreTrainedModel, config: CtxEncoderArguments):
super().__init__()
self.base_model = base_model = get_base_model(base_model)
self.exit_layer = config.layer_idx
base_model = get_base_model(base_model)
if "gte" in base_model.config.name_or_path:
self.base_model.encoder.layer = base_model.encoder.layer[: self.exit_layer]
base_model.encoder.layer = base_model.encoder.layer[: config.layer_idx]
else:
self.base_model.layers = base_model.layers[: self.exit_layer]
base_model.layers = base_model.layers[: config.layer_idx]
self.base_model = base_model
@property
def config(self):
@ -90,6 +91,8 @@ class EmbeddingOnly(nn.Module):
class PerLayerActivations(nn.Module):
def __init__(self, base_model: PreTrainedModel, config: CtxEncoderArguments):
super().__init__()
base_model = get_base_model(base_model) # remove lm head
base_model.layers = base_model.layers[:-1] # remove last attn block
self.base_model = base_model
@property
@ -103,7 +106,7 @@ class PerLayerActivations(nn.Module):
# Return all layers' activations except the last one
# from embeddings to the input of the last attn block
# Shape: (batch_size, num_layers, seq_len, hidden_size)
return torch.stack(outputs.hidden_states[:-1], dim=1)
return torch.stack(outputs.hidden_states, dim=1)
class CTX_ENCODER_TYPE(str, Enum):

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

@ -85,7 +85,7 @@ def apply_lora_to_layers(
) -> None:
layers = get_layers(model)
if position_ids is not None:
position_ids = position_ids.squeeze()
position_ids = position_ids.squeeze(0)
seq_lens = position_ids[torch.where(position_ids == 0)[0][1:] - 1]
seq_lens = torch.cat(
[seq_lens, torch.tensor([position_ids[-1]], device=seq_lens.device)]

80
src/ctx_to_lora/trainer.py
View file

@ -16,7 +16,7 @@ logger = logging.getLogger()
class DistillationTrainer(Trainer):
def __init__(self, *args, **kwargs):
self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
self.distillation_temperature = kwargs.pop("distillation_temperature", 1.0)
self.use_per_ctx_average_loss = kwargs.pop("use_per_ctx_average_loss", False)
super().__init__(*args, **kwargs)
def compute_loss(
@ -28,43 +28,23 @@ class DistillationTrainer(Trainer):
loss_kwargs["num_items_in_batch"] = num_items_in_batch
inputs = {**inputs, **loss_kwargs}
use_kl = "logprobs_vals" in inputs
# TODO: handle when val data doesnt have logprobs
if use_kl:
labels = inputs.pop("labels", None)
label_pos = torch.where(labels != -100)
# TODO: add efficient logits indexing?
# https://github.com/huggingface/trl/pull/2773
# + use `num_logits_to_keep` to limit the number of logits
outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True)
if not use_kl:
loss = outputs["loss"]
else:
target_logp = inputs.pop("logprobs_vals")
indices = inputs.pop("logprobs_indices")
# packed
if isinstance(target_logp, torch.Tensor):
target_logp = target_logp.squeeze(0)
indices = indices.squeeze(0)
# non-packed max_qas_per_sample=1
elif isinstance(target_logp[0], torch.Tensor):
target_logp = torch.cat([x[0] for x in target_logp], dim=0)
indices = torch.cat([x[0] for x in indices], dim=0)
# non-packed max_qas_per_sample>1
else:
# convert to tensor
raise NotImplementedError(
"Non-packed data with max_qas_per_sample > 1 is not supported yet."
)
target_logp = inputs.pop("logprobs_vals").squeeze(0)
indices = inputs.pop("logprobs_indices").squeeze(0)
assert label_pos[0].shape[0] == target_logp.shape[0], (
"Label positions and target log probabilities should have the same # tokens."
f"Got : {label_pos[0].shape[0]=} and {target_logp.shape[0]=}"
)
# TODO: per-ctx average loss?
##### KL loss
outputs_logits = outputs.logits[
label_pos[0], label_pos[1] - 1
@ -76,8 +56,48 @@ class DistillationTrainer(Trainer):
p = teacher_logp.exp()
logq = nn.functional.log_softmax(outputs_logits, dim=-1)
loss = -torch.sum(p * logq, dim=-1) * (self.distillation_temperature**2)
if reduction == "batchmean":
loss = -torch.sum(p * logq, dim=-1)
if self.use_per_ctx_average_loss:
n_queries_per_ctx = inputs["n_queries"].tolist()
position_ids = inputs["position_ids"].squeeze(0)
# account only label positions
label_mask = labels.squeeze(0) != -100
label_pos_ids = label_mask * position_ids
label_pos_ids_diff = label_pos_ids.diff(
append=torch.tensor([0], device=position_ids.device)
)
start_label_pos = torch.where((label_pos_ids_diff > 0) * ~label_mask)[0]
end_label_pos = torch.where((label_pos_ids_diff < 0) * label_mask)[0]
label_seq_lens = end_label_pos - start_label_pos
cu_label_seq_lens = torch.cumsum(label_seq_lens, dim=0)
start_indices = torch.cat(
(
torch.tensor([0], device=cu_label_seq_lens.device),
cu_label_seq_lens[:-1],
)
)
# these stack and split can be optimized but let's keep it simple
# Calculate mean loss for each QA pair
qa_losses = torch.stack(
[
loss[start:end].mean()
for start, end in zip(start_indices, cu_label_seq_lens)
]
)
# Split losses by context and calculate per-context average loss
per_ctx_losses = [
ctx.mean() for ctx in torch.split(qa_losses, n_queries_per_ctx)
]
# Final loss is the average of per-context losses
loss = torch.stack(per_ctx_losses).mean()
elif reduction == "batchmean":
loss = loss.mean()
elif reduction == "sum" and num_items_in_batch is not None:
loss = loss.sum() / num_items_in_batch
@ -85,8 +105,10 @@ class DistillationTrainer(Trainer):
##### unpack gen lora dict and compute regularization loss
l1_norm = 0
n_modules = len(gen_loras)
for module, lora in gen_loras.items():
l1_norm += lora["A"].abs().mean() + lora["B"].abs().mean()
l1_norm /= n_modules
# rough estimate of the losses (we only log the values from one step)
if (self.state.global_step == 1 and self.args.logging_first_step) or (
self.args.logging_strategy == IntervalStrategy.STEPS
@ -131,12 +153,8 @@ class ModulatedModelTrainer(Trainer):
if num_items_in_batch is not None:
loss_kwargs["num_items_in_batch"] = num_items_in_batch
inputs = {**inputs, **loss_kwargs}
# TODO: add efficient logits indexing?
# https://github.com/huggingface/trl/pull/2773
# + use `num_logits_to_keep` to limit the number of logits
outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True)
# Save past state if it exists
# TODO: this needs to be fixed and made cleaner later.
if self.args.past_index >= 0:
self._past = outputs[self.args.past_index]
@ -235,7 +253,11 @@ 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.")

20
intx_sft.py → train.py
View file

@ -239,6 +239,7 @@ def main():
training_args.gen_lora_l1_reg_coef = ctx_args.gen_lora_l1_reg_coef
training_args.use_kl_loss = ctx_args.use_kl_loss
training_args.use_per_ctx_average_loss = ctx_args.use_per_ctx_average_loss
if len([p for p in model.ctx_encoder.parameters() if p.requires_grad]):
raise ValueError("ctx_encoder contains trainable parameters")
@ -266,6 +267,9 @@ def main():
logger.info("Loading dataset...")
add_ctx_to_chat = not isinstance(model, ModulatedPretrainedModel)
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
_get_tokenized_dataset = partial(
get_tokenized_dataset,
@ -273,14 +277,17 @@ def main():
max_qas_per_sample=ctx_args.max_qas_per_sample,
base_model_max_len=model.base_model.config.max_position_embeddings,
tokenizer=tokenizer,
ctx_model_max_len=model.ctx_encoder.config.max_position_embeddings,
ctx_model_max_len=ctx_model_max_len,
ctx_tokenizer=ctx_tokenizer,
add_ctx_to_chat=add_ctx_to_chat,
use_kl_loss=ctx_args.use_kl_loss,
)
tokenized_ds = {"train": dict(), "validation": dict()}
splits = ["train"]
if training_args.eval_strategy != "no":
splits.append("validation")
tokenized_ds = {split: {} for split in splits}
for split, ds_names in zip(
["train", "validation"],
splits,
[data_args.train_ds_names, data_args.val_ds_names],
):
if not ds_names:
@ -323,6 +330,8 @@ def main():
train_ds_len = [len(ds) for ds in train_ds.values()]
total_len = sum(train_ds_len)
# TODO: cache this?
train_ds = interleave_datasets(
list(train_ds.values()),
probabilities=get_ds_prob(train_ds_len, total_len),
@ -339,7 +348,7 @@ def main():
ctx_args.max_packed_inp_len,
ctx_args.max_packed_ctx_len,
max_packed_size=-1,
num_proc=8,
num_proc=16,
)
logger.info(f"Train dataset length: {len(train_ds)}")
logger.info(
@ -370,7 +379,8 @@ def main():
logger.info("Compiling ctx_encoder_model")
ctx_base_model = model.ctx_encoder.base_model
compile_linear(ctx_base_model)
# model.ctx_encoder.compile(fullgraph=True, mode="max-autotune")
# TODO: why we can't compile ctx_base_model?
# ctx_base_model.compile(fullgraph=True, mode="max-autotune")
elif isinstance(model, PeftModel):
base_model = model.base_model

71
uv.lock generated
View file

@ -686,6 +686,8 @@ dependencies = [
{ name = "inflect" },
{ name = "jaxtyping" },
{ name = "jupyter" },
{ name = "kaggle" },
{ name = "kagglehub", extra = ["hf-datasets"] },
{ name = "liger-kernel" },
{ name = "matplotlib" },
{ name = "opt-einsum" },
@ -716,6 +718,8 @@ requires-dist = [
{ name = "inflect" },
{ name = "jaxtyping" },
{ name = "jupyter" },
{ name = "kaggle", specifier = ">=1.7.4.5" },
{ name = "kagglehub", extras = ["hf-datasets"], specifier = ">=0.3.12" },
{ name = "liger-kernel" },
{ name = "matplotlib" },
{ name = "opt-einsum", specifier = ">=3.4.0" },
@ -2109,6 +2113,52 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/43/6a/ca128561b22b60bd5a0c4ea26649e68c8556b82bc70a0c396eebc977fe86/jupyterlab_widgets-3.0.15-py3-none-any.whl", hash = "sha256:d59023d7d7ef71400d51e6fee9a88867f6e65e10a4201605d2d7f3e8f012a31c", size = 216571, upload-time = "2025-05-05T12:32:29.534Z" },
]
[[package]]
name = "kaggle"
version = "1.7.4.5"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "bleach" },
{ name = "certifi" },
{ name = "charset-normalizer" },
{ name = "idna" },
{ name = "protobuf" },
{ name = "python-dateutil" },
{ name = "python-slugify" },
{ name = "requests" },
{ name = "setuptools" },
{ name = "six" },
{ name = "text-unidecode" },
{ name = "tqdm" },
{ name = "urllib3" },
{ name = "webencodings" },
]
sdist = { url = "https://files.pythonhosted.org/packages/b1/02/b0c189a46531ea2b2691ae277508d2c80e5fd3d757083283c5cc27800ca8/kaggle-1.7.4.5.tar.gz", hash = "sha256:1d9821bd6a6a1470741c76d26495a18475b5a7bfe0c80b19191254b2735d41dd", size = 336100, upload-time = "2025-05-08T21:17:20.081Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/14/83/7f29c7abe0d5dc769dad7da993382c3e4239ad63e1dd58414d129e0a4da2/kaggle-1.7.4.5-py3-none-any.whl", hash = "sha256:9732afb1c073f14acc7e49dfab98456f887d10b735bcd6348bc1340e92393882", size = 181238, upload-time = "2025-05-08T21:17:18.245Z" },
]
[[package]]
name = "kagglehub"
version = "0.3.12"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "packaging" },
{ name = "pyyaml" },
{ name = "requests" },
{ name = "tqdm" },
]
sdist = { url = "https://files.pythonhosted.org/packages/ee/75/c48635ebfb8977945e1a3bf19ec78fba7e3c67b0e808eb633eb9500602ca/kagglehub-0.3.12.tar.gz", hash = "sha256:45e75854630a30605b794eb786b3757beccbbea1acca71600642f67b60e0d7bf", size = 112292, upload-time = "2025-04-23T23:04:57.077Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/49/bf/c2a24567bb6bd80c1fe7cb2ed1a332666476f69c313256aff96094bef93e/kagglehub-0.3.12-py3-none-any.whl", hash = "sha256:6eeb7c5c8e1f25a28e9b26d3d8ed0c24a4a32f763ea5de6e1b3aabba442f2a26", size = 67996, upload-time = "2025-04-23T23:04:55.694Z" },
]
[package.optional-dependencies]
hf-datasets = [
{ name = "datasets" },
{ name = "pandas" },
]
[[package]]
name = "kiwisolver"
version = "1.4.8"
@ -3952,6 +4002,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/45/58/38b5afbc1a800eeea951b9285d3912613f2603bdf897a4ab0f4bd7f405fc/python_multipart-0.0.20-py3-none-any.whl", hash = "sha256:8a62d3a8335e06589fe01f2a3e178cdcc632f3fbe0d492ad9ee0ec35aab1f104", size = 24546, upload-time = "2024-12-16T19:45:44.423Z" },
]
[[package]]
name = "python-slugify"
version = "8.0.4"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "text-unidecode" },
]
sdist = { url = "https://files.pythonhosted.org/packages/87/c7/5e1547c44e31da50a460df93af11a535ace568ef89d7a811069ead340c4a/python-slugify-8.0.4.tar.gz", hash = "sha256:59202371d1d05b54a9e7720c5e038f928f45daaffe41dd10822f3907b937c856", size = 10921, upload-time = "2024-02-08T18:32:45.488Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/a4/62/02da182e544a51a5c3ccf4b03ab79df279f9c60c5e82d5e8bec7ca26ac11/python_slugify-8.0.4-py2.py3-none-any.whl", hash = "sha256:276540b79961052b66b7d116620b36518847f52d5fd9e3a70164fc8c50faa6b8", size = 10051, upload-time = "2024-02-08T18:32:43.911Z" },
]
[[package]]
name = "pytz"
version = "2025.2"
@ -4763,6 +4825,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/6a/9e/2064975477fdc887e47ad42157e214526dcad8f317a948dee17e1659a62f/terminado-0.18.1-py3-none-any.whl", hash = "sha256:a4468e1b37bb318f8a86514f65814e1afc977cf29b3992a4500d9dd305dcceb0", size = 14154, upload-time = "2024-03-12T14:34:36.569Z" },
]
[[package]]
name = "text-unidecode"
version = "1.3"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/ab/e2/e9a00f0ccb71718418230718b3d900e71a5d16e701a3dae079a21e9cd8f8/text-unidecode-1.3.tar.gz", hash = "sha256:bad6603bb14d279193107714b288be206cac565dfa49aa5b105294dd5c4aab93", size = 76885, upload-time = "2019-08-30T21:36:45.405Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/a6/a5/c0b6468d3824fe3fde30dbb5e1f687b291608f9473681bbf7dabbf5a87d7/text_unidecode-1.3-py2.py3-none-any.whl", hash = "sha256:1311f10e8b895935241623731c2ba64f4c455287888b18189350b67134a822e8", size = 78154, upload-time = "2019-08-30T21:37:03.543Z" },
]
[[package]]
name = "tiktoken"
version = "0.9.0"

24
watcher.py
View file

@ -31,9 +31,14 @@ class Watcher:
def watch(self):
self.files = self.get_files()
new_files = self.files - self.last_files
self.last_files = self.files
return sorted(list(new_files))
def update(self, file):
if file in self.last_files:
return
self.last_files.add(file)
print(f"Added {file} to evaluated files.")
def save_state(self):
with open("watcher_state.yaml", "w") as f:
yaml.dump({"last_files": self.last_files}, f)
@ -79,14 +84,15 @@ if __name__ == "__main__":
# TODO: have to change this for bigger models
eval_batch_size = 8
eval_batch_size_gen = 32
eval_batch_size_gen = 16
try:
metrics = run_eval(
checkpoint_path=file,
eval_batch_size=eval_batch_size,
split="validation",
generative=False,
)
# metrics = run_eval(
# checkpoint_path=file,
# eval_batch_size=eval_batch_size,
# split="validation",
# generative=False,
# )
metrics = {}
except FileNotFoundError as e:
print(f"Error evaluating {file}: {e}. The checkpoint might be deleted.")
continue
@ -100,6 +106,7 @@ if __name__ == "__main__":
except FileNotFoundError as e:
print(f"Error evaluating {file}: {e}. The checkpoint might be deleted.")
gen_metrics = {}
file = ""
metrics.update(gen_metrics)
for k in metrics:
wandb.log(metrics[k], step=curstep)
@ -107,5 +114,6 @@ if __name__ == "__main__":
print(f"Logged metrics: {metrics}")
print("=" * 80)
clear_gpu()
watcher.update(file)
watcher.save_state()