iclr cleanup

scripts/large_exp/gemma_qa_exp_4_gpus.sh

@@ -1,30 +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
-
-port=$((10000 + ($SLURM_JOBID % 50000)))
-echo "Using port: $port"
-
-# 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"
-    "--num_train_epochs=1"
-    "--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 --config_file accelerate_config.yaml --main_process_port $port \
---num_processes=4 --gpu_ids all train.py $1 "${default_args[@]}" "${@:2}"

scripts/large_exp/gemma_qa_exp_8_gpus.sh

@@ -1,30 +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
-
-port=$((10000 + ($SLURM_JOBID % 50000)))
-echo "Using port: $port"
-
-# 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"
-    "--num_train_epochs=1"
-    "--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 --config_file accelerate_config.yaml --main_process_port $port \
---num_processes=8 --gpu_ids all train.py $1 "${default_args[@]}" "${@:2}"

scripts/main_exp/eval/base_model.sh

@@ -0,0 +1,8 @@
+# no truncation
+WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1
+
+# w/ truncation
+WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --truncate_if_too_long_inp
+
+# no context
+WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --remove_context

scripts/main_exp/eval/cd.sh

@@ -0,0 +1,6 @@
+# qa
+uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=4
+
+
+# longbench
+uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=1

scripts/main_exp/eval/cd_oracle.sh

@@ -0,0 +1,6 @@
+# qa
+uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp
+
+
+# longbench
+uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp

scripts/main_exp/eval/d2l.sh

@@ -0,0 +1,13 @@
+# main results
+# batched
+WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1
+
+# iterative
+WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1 --use_iterative_mode
+
+# query internalization
+WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad --split test --eval_batch_size_gen=1 --flip_ctx_inp
+
+# replaced squad context
+WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad_assistant_ctx_no_passage --split test
+WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad_negative_no_passage --split test

scripts/main_exp/eval/llmlingua.sh

@@ -0,0 +1,5 @@
+for dataset in squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e; do
+  for rate in 0.9 0.8 0.6 0.4 0.2 0.1; do
+    WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets "$dataset" --split test --eval_batch_size_gen=1 --use_llmlingua --llmlingua_compression_rate "$rate" --truncate_if_too_long_ctx
+  done
+done

scripts/main_exp/eval/t2l.sh

@@ -0,0 +1,4 @@
+# download t2l checkpoint
+uv run huggingface-cli download SakanaAI/text-to-lora --local-dir . --include "trained_t2l/gemma_2b_t2l"
+
+WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen=1 --use_t2l

scripts/main_exp/gen_data.sh

@@ -0,0 +1,20 @@
+# download fineweb_edu to `data/raw_datasets/fineweb_edu
+uv run data/download_fineweb_edu.py
+
+# generate qa data
+# run from 000 to 013
+for shard_id in $(seq -f "%03g" 0 1); do
+  uv run data/generate_fw_edu_qa_v2.py --shard_pattern "${shard_id}_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048;
+  uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/${shard_id}*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
+
+  # self-generated response QA data
+  uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern "data/raw_datasets/fw_qa_v2/min_0_to_2000/${shard_id}*_level_1*" --closed_qa_prob 1.0
+done
+
+
+# val split
+uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_v2/min_0_to_2000/*_level_0_val.parquet'
+
+# self-gen data for other ds 
+uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names squad_compact ropes_compact drop_compact --closed_qa_prob 1.0
+uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names pwc_compact --closed_qa_prob 0.0

scripts/main_exp/train.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+
+port=29051
+
+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"
+)
+
+uv run accelerate launch --config_file accelerate_config.yaml --main_process_port $port \
+--num_processes=8 --gpu_ids all train.py \
+configs/main_exp/self_gen_lv1_closed_qa_1_l2l.yaml \
+"${default_args[@]}" \
+--model_name_or_path=google/gemma-2-2b-it \
+--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 \
+--max_steps=80000 --gradient_accumulation_steps=8 --max_packed_inp_len=4096 \
+--max_packed_ctx_len=4096 --use_per_ctx_average_loss=True --use_kl_loss=True \
+--quantize_ctx_encoder=True 

scripts/niah/0-gen_data.sh

@@ -0,0 +1 @@
+uv run data/generate_ctx_magic_number.py

scripts/niah/1-train.sh

@@ -0,0 +1,22 @@
+WANDB_MODE=disabled uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml \
+--model_name_or_path=google/gemma-2-2b-it \
+--num_train_epochs=1 \
+--per_device_train_batch_size=-1 \
+--gradient_accumulation_steps=32 \
+--per_device_eval_batch_size=16 \
+--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 --logging_steps=10 \
+--save_steps=1000 --learning_rate=4e-5 \
+--lora_dropout=0.0 --neftune_noise_alpha=0 \
+--per_rank_gen=True --per_layer_processing=True \
+--gen_lora_l1_reg_coef=1 --use_sequence_packing=True \
+--max_packed_inp_len=2048 --max_packed_ctx_len=2048 \
+--dataloader_num_workers=0 --dataloader_prefetch_factor=None \
+--ctx_encoder_type=early_exit --n_latent_queries=208 \
+--eval_on_start=True --lora_r=8 \
+--max_ctx_chunk_len=512 --min_ctx_chunk_len=25 \
+--num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' \
+--max_val_samples_per_ds=100 --seed=1 --use_per_ctx_average_loss=True

scripts/niah/2-eval.sh

@@ -0,0 +1 @@
+WANDB_MODE=disabled uv run run_eval.py --checkpoint_path CHECKPOINT_PATH --datasets ctx_magic_number_32_1024 ctx_magic_number_1024_2048 ctx_magic_number_2048_3072 ctx_magic_number_3072_4096 ctx_magic_number_4096_5120 ctx_magic_number_5120_6144 ctx_magic_number_6144_7168 ctx_magic_number_7168_8192 ctx_magic_number_8192_9216 ctx_magic_number_9216_10240 ctx_magic_number_10240_11264 ctx_magic_number_11264_12288 ctx_magic_number_12288_13312 ctx_magic_number_13312_14336 ctx_magic_number_14336_15360 ctx_magic_number_15360_16384 ctx_magic_number_16384_20480 ctx_magic_number_20480_24576 ctx_magic_number_24576_28672 ctx_magic_number_28672_32768 ctx_magic_number_32768_40960 ctx_magic_number_40960_49152 ctx_magic_number_49152_57344 ctx_magic_number_57344_65536 ctx_magic_number_65536_73728 ctx_magic_number_73728_81920 ctx_magic_number_81920_90112 ctx_magic_number_90112_98304 ctx_magic_number_98304_106496 ctx_magic_number_106496_114688 ctx_magic_number_114688_122880 ctx_magic_number_122880_131072 --max_ctx_chunk_len=1024 --split test --eval_batch_size_gen=4

scripts/sbatch_1_gpu.sh

@@ -1,10 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=ctxlora
-#SBATCH --nodes=1
-#SBATCH --partition=sakura-gpu
-#SBATCH --gpus=1
-#SBATCH --output=slurm_logs/%x-%j.out
-#SBATCH --error=slurm_logs/%x-%j.out
-
-
-"$@"

scripts/short_ctx/gemma_qa_short_ctx_exp_default.sh

@@ -1,29 +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
-
-port=$((10000 + ($SLURM_JOBID % 50000)))
-echo "Using port: $port"
-
-# 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 --config_file accelerate_config.yaml --main_process_port $port \
---num_processes=4 --gpu_ids all train.py $1 "${default_args[@]}" "${@:2}"

scripts/short_ctx/gemma_qa_short_ctx_exp_default_8_gpus.sh

@@ -1,29 +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
-
-port=$((10000 + ($SLURM_JOBID % 50000)))
-echo "Using port: $port"
-
-# 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 --config_file accelerate_config.yaml --main_process_port $port \
---num_processes=8 --gpu_ids all train.py $1 "${default_args[@]}" "${@:2}"

scripts/submit_tiny_exp_jobs.sh

@@ -1,48 +0,0 @@
-#!/bin/bash
-
-# Script to submit SLURM jobs for all configs in configs/tiny_exp/
-# with random sleep intervals between submissions
-
-CONFIG_DIR="configs/tiny_exp"
-SCRIPT_PATH="scripts/short_ctx/gemma_qa_short_ctx_exp_default.sh"
-
-# Check if config directory exists
-if [ ! -d "$CONFIG_DIR" ]; then
-    echo "Error: Config directory $CONFIG_DIR does not exist"
-    exit 1
-fi
-
-# Check if script exists
-if [ ! -f "$SCRIPT_PATH" ]; then
-    echo "Error: Script $SCRIPT_PATH does not exist"
-    exit 1
-fi
-
-# Get all yaml config files
-configs=($(find "$CONFIG_DIR" -name "*.yaml" -type f))
-
-if [ ${#configs[@]} -eq 0 ]; then
-    echo "No .yaml config files found in $CONFIG_DIR"
-    exit 1
-fi
-
-echo "Found ${#configs[@]} config files:"
-for config in "${configs[@]}"; do
-    echo "  - $config"
-done
-
-echo ""
-echo "Starting job submissions..."
-
-# Submit jobs with random sleep intervals
-for config in "${configs[@]}"; do
-    echo "Submitting job for config: $config"
-    sbatch "$SCRIPT_PATH" "$config"
-
-    # Random sleep between 3-30 seconds
-    sleep_time=$((3 + RANDOM % 28))
-    echo "Waiting $sleep_time seconds before next submission..."
-    sleep $sleep_time
-done
-
-echo "All jobs submitted!"

scripts/toy_exp.sh

@@ -1,17 +0,0 @@
-# data gen
-uv run data/generate_ctx_magic_num.py
-
-# train
-# WANDB_PROJECT=ctx-magic-num srun --partition=aiscilow --gpus=1 --unbuffered uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=32 --per_device_eval_batch_size=16 --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=0 --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=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=-1 --max_val_samples_per_ds=100 --seed=1
-# WANDB_PROJECT=ctx-magic-num srun --partition=aiscilow --gpus=1 --unbuffered uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=32 --per_device_eval_batch_size=16 --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=0 --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=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=512 --min_ctx_chunk_len=25 --num_chunk_probs='{"1":"0.5", "2":"0.5"}' --max_val_samples_per_ds=100 --seed=1
-# WANDB_PROJECT=ctx-magic-num srun --partition=aiscilow --gpus=1 --unbuffered uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=32 --per_device_eval_batch_size=16 --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=0 --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=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=512 --min_ctx_chunk_len=25 --num_chunk_probs='{"1":"0.5", "2":"0.25", "3":"0.125", "4":"0.125"}' --max_val_samples_per_ds=100 --seed=1
-WANDB_PROJECT=ctx-magic-num run uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=16 --per_device_eval_batch_size=16 --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 --logging_steps=10 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=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 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=512 --min_ctx_chunk_len=25 --num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' --max_val_samples_per_ds=100 --seed=1 --use_per_ctx_average_loss=True
-
-# self-distill version
-# ce
-# WANDB_PROJECT=ctx-magic-num run uv run train.py configs/toy_exp/ctx_magic_number_32_256_self_gen.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=16 --per_device_eval_batch_size=16 --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 --logging_steps=10 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=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=per_layer_activations --n_latent_queries=8 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=-1 --max_val_samples_per_ds=100 --seed=1 --notes="1 l1norm + ce loss + 1 chunk" --use_per_ctx_average_loss=False
-# kl
-# WANDB_PROJECT=ctx-magic-num run uv run train.py configs/toy_exp/ctx_magic_number_32_256_self_gen.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=16 --per_device_eval_batch_size=16 --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 --logging_steps=10 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=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=per_layer_activations --n_latent_queries=8 --use_kl_loss=True --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=-1 --max_val_samples_per_ds=100 --seed=1 --notes="1 l1norm + kl loss + 1 chunk" --use_per_ctx_average_loss=False
-
-# eval
-WANDB_MODE=disabled srun --partition=aiscilow --gpus=1 --unbuffered uv run run_eval.py --checkpoint_path CHECKPOINT_PATH --datasets ctx_magic_number_32_1024 ctx_magic_number_1024_2048 ctx_magic_number_2048_3072 ctx_magic_number_3072_4096 ctx_magic_number_4096_5120 ctx_magic_number_5120_6144 ctx_magic_number_6144_7168 ctx_magic_number_7168_8192 ctx_magic_number_8192_9216 ctx_magic_number_9216_10240 ctx_magic_number_10240_11264 ctx_magic_number_11264_12288 ctx_magic_number_12288_13312 ctx_magic_number_13312_14336 ctx_magic_number_14336_15360 ctx_magic_number_15360_16384 ctx_magic_number_16384_20480 ctx_magic_number_20480_24576 ctx_magic_number_24576_28672 ctx_magic_number_28672_32768 ctx_magic_number_32768_40960 ctx_magic_number_40960_49152 ctx_magic_number_49152_57344 ctx_magic_number_57344_65536 ctx_magic_number_65536_73728 ctx_magic_number_73728_81920 ctx_magic_number_81920_90112 ctx_magic_number_90112_98304 ctx_magic_number_98304_106496 ctx_magic_number_106496_114688 ctx_magic_number_114688_122880 ctx_magic_number_122880_131072 --max_ctx_chunk_len=1024 --split test --eval_batch_size_gen=4