diff --git a/configs/pretrain_all_xl_3.yaml b/configs/pretrain_all_xl_3.yaml new file mode 100644 index 0000000..e19f84a --- /dev/null +++ b/configs/pretrain_all_xl_3.yaml @@ -0,0 +1,59 @@ +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: 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: 5 +weight_decay: 0.01 +# +warmup_steps: 100 + +dataloader_prefetch_factor: 8 +dataloader_num_workers: 8 +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj +# data +train_ds_names: + - fw_qa_3_mini + - ctx_qa + - pwc + - hotpot_qa + - squad + - drop + - narrativeqa + - quoref + - ropes + - synthetic_convqa + +val_ds_names: + - fw_qa_xl + - ctx_qa + - pwc + - hotpot_qa + - squad + +load_best_model_at_end: true +metric_for_best_model: eval_pwc_loss diff --git a/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2_per_layer_processing_4e-5_lr.sh b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2_per_layer_processing_4e-5_lr.sh new file mode 100644 index 0000000..fa2678b --- /dev/null +++ b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2_per_layer_processing_4e-5_lr.sh @@ -0,0 +1,34 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --exclude=slurm0-a3nodeset-2 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +# module load +# module load cuda/12.1 +# module load cudnn/8.9.7 +# module load nccl/cuda-12.1/2.18.3 +# module load hpcx/2.20 + +# export OMP_NUM_THREADS=24 +# export TRITON_CACHE_DIR=/tmp/.triton/ +. ~/miniconda3/etc/profile.d/conda.sh +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora +# eval "$@" + +accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29562 intx_sft.py configs/pretrain_all_xl.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --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 diff --git a/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2_per_layer_processing_decoder_2.sh b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2_per_layer_processing_decoder_2.sh new file mode 100644 index 0000000..685fd4d --- /dev/null +++ b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2_per_layer_processing_decoder_2.sh @@ -0,0 +1,35 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --exclude=slurm0-a3nodeset-2 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +# module load +# module load cuda/12.1 +# module load cudnn/8.9.7 +# module load nccl/cuda-12.1/2.18.3 +# module load hpcx/2.20 + +# export OMP_NUM_THREADS=24 +# export TRITON_CACHE_DIR=/tmp/.triton/ +. ~/miniconda3/etc/profile.d/conda.sh +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora +# eval "$@" + +accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29563 intx_sft.py configs/pretrain_all_xl.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --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 \ +--decoder_depth=2 \ +--lora_r=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-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 diff --git a/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2_per_layer_processing_r16.sh b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2_per_layer_processing_r16.sh new file mode 100644 index 0000000..f039792 --- /dev/null +++ b/scripts/ablate_arch/gemma_xl_per_rank_gen_fac2_per_layer_processing_r16.sh @@ -0,0 +1,34 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --exclude=slurm0-a3nodeset-2 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +# module load +# module load cuda/12.1 +# module load cudnn/8.9.7 +# module load nccl/cuda-12.1/2.18.3 +# module load hpcx/2.20 + +# export OMP_NUM_THREADS=24 +# export TRITON_CACHE_DIR=/tmp/.triton/ +. ~/miniconda3/etc/profile.d/conda.sh +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora +# eval "$@" + +accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29562 intx_sft.py configs/pretrain_all_xl.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --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=16 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-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 diff --git a/scripts/fw_qa_3/gemma_per_rank_fac2_per_layer.sh b/scripts/fw_qa_3/gemma_per_rank_fac2_per_layer.sh new file mode 100644 index 0000000..123bdf8 --- /dev/null +++ b/scripts/fw_qa_3/gemma_per_rank_fac2_per_layer.sh @@ -0,0 +1,34 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#SBATCH --exclude=slurm0-a3nodeset-2 +#SBATCH --gpus=4 +#SBATCH --output=outputs/%x-%j.out +#SBATCH --error=outputs/%x-%j.out + +# module load +# module load cuda/12.1 +# module load cudnn/8.9.7 +# module load nccl/cuda-12.1/2.18.3 +# module load hpcx/2.20 + +# export OMP_NUM_THREADS=24 +# export TRITON_CACHE_DIR=/tmp/.triton/ +. ~/miniconda3/etc/profile.d/conda.sh +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora +# eval "$@" + +accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29562 intx_sft.py configs/pretrain_all_xl_3.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --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 diff --git a/src/ctx_to_lora/data_utils.py b/src/ctx_to_lora/data_utils.py index 3ba1da9..8dc7b50 100644 --- a/src/ctx_to_lora/data_utils.py +++ b/src/ctx_to_lora/data_utils.py @@ -98,6 +98,13 @@ DS_KWARGS = { split="train", ), ), + "fw_qa_3_mini": dict( + train=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_3/*[!val].parquet"), + split="train[:100000]", + ), + ), "ctx_qa": dict( train=dict( path="parquet", @@ -250,7 +257,7 @@ def add_repeat_prompt_fn(samples): ): # Only process if the context is not already in the set if ctx in unique_contexts: - continue + continue # remove too long contexts if len(ctx) > 1000: continue diff --git a/src/ctx_to_lora/modeling_utils.py b/src/ctx_to_lora/modeling_utils.py index 3ccd908..55e22ab 100644 --- a/src/ctx_to_lora/modeling_utils.py +++ b/src/ctx_to_lora/modeling_utils.py @@ -756,14 +756,14 @@ class HyperLoRA(nn.Module): if self.config.per_layer_processing: layers = [ # nn.Linear(self.d_latent, self.d_latent), - Mix( - "bs n_layers n_modules r d -> bs n_layers n_modules r d0", - weight_shape="n_layers d d0", - bias_shape="n_layers d0", - n_layers=self.n_layers, - d=self.d_latent, - d0=self.d_latent, - ), + # Mix( + # "bs n_layers n_modules r d -> bs n_layers n_modules r d0", + # weight_shape="n_layers d d0", + # bias_shape="n_layers d0", + # n_layers=self.n_layers, + # d=self.d_latent, + # d0=self.d_latent, + # ), ResMLPBlockPerLayer( self.n_layers, self.d_latent, @@ -772,27 +772,29 @@ class HyperLoRA(nn.Module): # self.num_modules, # self.lora_config.r, # self.d_latent, - ), + ) + for _ in range(4) ] else: layers = [ - nn.Linear(self.d_latent, self.d_latent), + # nn.Linear(self.d_latent, self.d_latent), MLPResidualBlock( input_size=self.config.latent_size, hidden_size=self.config.latent_size * 4, output_size=self.config.latent_size, dropout_rate=getattr(self.config, "dropout_rate", 0), - ), - ] - if self.config.use_token_mixing: - layers = [ - ResMLPTokenMixingPerLayer( - self.n_layers, - self.num_modules, - self.lora_config.r, - self.d_latent, ) - ] + layers + for _ in range(4) + ] + # if self.config.use_token_mixing: + # layers = [ + # ResMLPTokenMixingPerLayer( + # self.n_layers, + # self.num_modules, + # self.lora_config.r, + # self.d_latent, + # ) + # ] + layers self.layers = nn.Sequential(*layers)