diff --git a/batch.sh b/batch.sh deleted file mode 100755 index 9e1a41d..0000000 --- a/batch.sh +++ /dev/null @@ -1,29 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=ctx_to_lora -#SBATCH --partition=a3 -#SBATCH --gpus=4 -#SBATCH --exclude=slurm0-a3nodeset-[0-5] -#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/llm-modulator -# eval "$@" - -accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ ---gpu_ids all --main_process_port 29551 intx_sft.py configs/fw_qa_large_ctx_pwc_hotpot_squad.yaml \ ---model_name_or_path=meta-llama/Llama-3.1-8B-Instruct --num_train_epochs=5.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,up_proj,gate_proj --num_blocks=1 --num_self_attends_per_block=32 \ ---self_attention_widening_factor=4 --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 --resume_from_checkpoint=train_outputs/runs/Jan18_14-07-58_slurm0-a3nodeset-6_3de73fea/checkpoint-30000 diff --git a/batch_gemma_2_2b.sh b/batch_gemma_2_2b.sh deleted file mode 100755 index dcc8377..0000000 --- a/batch_gemma_2_2b.sh +++ /dev/null @@ -1,30 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=ctx_to_lora -#SBATCH --partition=a3 -#SBATCH --gpus=8 -#SBATCH --exclude=slurm0-a3nodeset-[0-5] -#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/llm-modulator -# eval "$@" - -accelerate launch --num_processes=8 --gradient_accumulation_steps=4 --gradient_clipping=1.0 \ ---gpu_ids all --main_process_port 29553 intx_sft.py configs/fw_qa_large_ctx_pwc_hotpot_squad.yaml \ ---model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.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,up_proj,gate_proj --num_blocks=1 --num_self_attends_per_block=32 \ ---self_attention_widening_factor=4 --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 -# --resume_from_checkpoint=train_outputs/runs/Jan18_14-22-38_slurm0-a3nodeset-10_c9340910/checkpoint-15000/ diff --git a/batch_gemma_llama_new_up_and_down.sh b/batch_gemma_llama_new_up_and_down.sh deleted file mode 100644 index 3f80919..0000000 --- a/batch_gemma_llama_new_up_and_down.sh +++ /dev/null @@ -1,31 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=gemma_llama_instruct -#SBATCH --partition=a3 -#SBATCH --nodes=1 -#SBATCH --gpus=8 -#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/llm-modulator -# eval "$@" - -accelerate launch --num_processes=8 --gradient_accumulation_steps=1 --gradient_clipping=1.0 \ ---gpu_ids all --main_process_port 29554 intx_sft.py configs/pretrain_all.yaml \ ---model_name_or_path=google/gemma-2-2b-it --num_train_epochs=10.1 --per_device_train_batch_size=32 \ ---gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ ---target_modules=up_proj,down_proj --num_blocks=1 --num_self_attends_per_block=16 \ ---eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 \ ---neftune_noise_alpha=1 --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.1-8B-Instruct \ ---use_sequence_packing=True - diff --git a/batch_gemma_llama_new_up_and_down_ln.sh b/batch_gemma_llama_new_up_and_down_ln.sh deleted file mode 100644 index 4fb353d..0000000 --- a/batch_gemma_llama_new_up_and_down_ln.sh +++ /dev/null @@ -1,31 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=gemma_llama_instruct -#SBATCH --partition=a3 -#SBATCH --nodes=1 -#SBATCH --gpus=8 -#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/llm-modulator -# eval "$@" - -accelerate launch --num_processes=8 --gradient_accumulation_steps=1 --gradient_clipping=1.0 \ ---gpu_ids all --main_process_port 29554 intx_sft.py configs/pretrain_all.yaml \ ---model_name_or_path=google/gemma-2-2b-it --num_train_epochs=10.1 --per_device_train_batch_size=32 \ ---gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ ---target_modules=up_proj,down_proj --extra_modules=input_layernorm,post_attention_layernorm --num_blocks=1 --num_self_attends_per_block=16 \ ---eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 \ ---neftune_noise_alpha=1 --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.1-8B-Instruct \ ---use_sequence_packing=True - diff --git a/batch_gemma_llama_vision.sh b/batch_gemma_llama_vision.sh deleted file mode 100755 index 03ff5c8..0000000 --- a/batch_gemma_llama_vision.sh +++ /dev/null @@ -1,30 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=gemma_llama_vision -#SBATCH --partition=a3 -#SBATCH --nodes=1 -#SBATCH --gpus=8 -#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/llm-modulator -# eval "$@" - -accelerate launch --num_processes=8 --gradient_accumulation_steps=4 --gradient_clipping=1.0 \ ---gpu_ids all --main_process_port 29554 intx_sft.py configs/pretrain_all.yaml \ ---model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=4 \ ---gradient_accumulation_steps=4 --per_device_eval_batch_size=4 --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 - diff --git a/batch_llama_2_7b.sh b/batch_llama_2_7b.sh deleted file mode 100755 index 76ba443..0000000 --- a/batch_llama_2_7b.sh +++ /dev/null @@ -1,29 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=ctx_to_lora -#SBATCH --partition=a3 -#SBATCH --gpus=4 -#SBATCH --exclude=slurm0-a3nodeset-[0-5] -#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/llm-modulator -# eval "$@" - -accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \ ---gpu_ids all --main_process_port 29552 intx_sft.py configs/fw_qa_large_ctx_pwc_hotpot_squad.yaml \ ---model_name_or_path=meta-llama/Llama-2-7b-chat-hf --num_train_epochs=5.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,up_proj,gate_proj --num_blocks=1 --num_self_attends_per_block=32 \ ---self_attention_widening_factor=4 --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 diff --git a/batch_llama_2b_.sh b/batch_llama_2b_.sh deleted file mode 100755 index a2bf1f6..0000000 --- a/batch_llama_2b_.sh +++ /dev/null @@ -1,30 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=gemma_pixtral -#SBATCH --partition=a3 -#SBATCH --nodes=1 -#SBATCH --gpus=8 -#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/llm-modulator -# eval "$@" - -accelerate launch --num_processes=8 --gradient_accumulation_steps=4 --gradient_clipping=1.0 \ ---gpu_ids all --main_process_port 29554 intx_sft.py configs/pretrain_all.yaml \ ---model_name_or_path=meta-llama/Llama-3.2.3B-Instruct --num_train_epochs=5.1 --per_device_train_batch_size=4 \ ---gradient_accumulation_steps=4 --per_device_eval_batch_size=4 --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=1024 \ ---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=AIDC-AI/Ovis1.6-Gemma2-9B - diff --git a/batch_llama_3_8b_1k_latent.sh b/batch_llama_3_8b_1k_latent.sh deleted file mode 100755 index 05bdd3a..0000000 --- a/batch_llama_3_8b_1k_latent.sh +++ /dev/null @@ -1,29 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=ctx_to_lora -#SBATCH --partition=a3 -#SBATCH --nodes=1 -#SBATCH --gpus=8 -#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/llm-modulator -# eval "$@" - -accelerate launch --num_processes=8 --gradient_accumulation_steps=4 --gradient_clipping=1.0 \ ---gpu_ids all --main_process_port 29555 intx_sft.py configs/pretrain_all.yaml \ ---model_name_or_path=meta-llama/Llama-3.1-8B-Instruct --num_train_epochs=10.1 --per_device_train_batch_size=4 \ ---gradient_accumulation_steps=4 --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=32 \ ---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=1024 \ ---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=Alibaba-NLP/gte-large-en-v1.5 diff --git a/configs/pretrain_all_2.yaml b/configs/pretrain_all_2.yaml new file mode 100644 index 0000000..9b4ae2a --- /dev/null +++ b/configs/pretrain_all_2.yaml @@ -0,0 +1,61 @@ +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.00002 +# 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.05 +target_modules: + - down_proj + - up_proj + +# data +train_ds_names: +- fw_qa_2 +- 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/dl_fineweb_edu.py b/dl_fineweb_edu.py new file mode 100644 index 0000000..70e09f3 --- /dev/null +++ b/dl_fineweb_edu.py @@ -0,0 +1,39 @@ +import random +from glob import glob + +from datasets import load_dataset, Dataset +from huggingface_hub import snapshot_download +from transformers import set_seed + + +def remove_too_long(samples): + return [len(text) < 10_000 for text in samples["text"]] + + +if __name__ == "__main__": + set_seed(42) + fw_dir = "./data/raw_datasets/fineweb_edu/" + snapshot_download( + "HuggingFaceFW/fineweb-edu", + repo_type="dataset", + local_dir=fw_dir, + allow_patterns="sample/100BT/*", + ) + # # https://github.com/huggingface/datasets/issues/7047#issuecomment-2233163406 + # num_shards_per_file = 16 + # sharded_fw_dir = "./data/raw_datasets/fineweb_sharded/" + # output_path_template = f"{sharded_fw_dir}" + "/{i:02d}_{idx:05d}.parquet" + + # for i, f in enumerate(sorted(glob(f"{fw_dir}/sample/100BT/*.parquet"))): + # # ~1M rows ~= 2GB mem required per file + # ds = Dataset.from_parquet(f) + # ds = ds.filter(remove_too_long, batched=True) + # print(f"Filtered ds size: {len(ds)}") + # ds = ds.shuffle(seed=42 + i) + # # take one shard (from 16 shards) per file + # # idx = random.sample(range(num_shards_per_file), 1)[0] + # # shard = ds.shard(index=idx, num_shards=num_shards_per_file, contiguous=False) + # # shard.to_parquet(output_path_template.format(index=i)) + # for idx in range(num_shards_per_file): + # shard = ds.shard(index=idx, num_shards=num_shards_per_file, contiguous=True) + # shard.to_parquet(output_path_template.format(i=i, idx=idx)) diff --git a/generate_fw_edu_qa_vllm.py b/generate_fw_edu_qa_vllm.py new file mode 100644 index 0000000..4b536a3 --- /dev/null +++ b/generate_fw_edu_qa_vllm.py @@ -0,0 +1,232 @@ +import os +import sys +import random +import re +from glob import glob + +import pandas as pd +from datasets import Dataset +from vllm import LLM, SamplingParams +from tqdm import tqdm +from datasets import load_dataset +from openai import OpenAI + +SYSTEM_TEMPLATE = ( + "You are a creative and helpful assistant.\n" + "You are tasked with generating questions for reading comprehension tests.\n" + "You will be given a context and you need to generate questions and corresponding answers from the given context.\n" + "The questions should be highly specific to the information provided in the context, not general questions that suits any context.\n" + "**DO NOT** halucinate or make up information." +) + +# based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811) +PROMPT_TEMPLATE = ( + "### Instructions ###\n" + "Generate questions and corresponding answers from the given context. The questions should be highly specific to the " + "information provided in the context, not general questions that suits any context.\n\n\n" + "### Context ###\n" + "{context}" + "\n\n\n" + "### Additional Instructions ###\n" + "Rules to follow when generate the questions:\n" + "1. The questions must be specific to the given context and fully answerable from information present in given context.\n" + "2. Make sure the questions are clear and unambiguous.\n" + "3. Phrases like 'based on the provided context', 'according to the context', etc, are **NOT ALLOWED** to appear in " + "the questions.\n" + "4. The questions should not overlap. They should be diverse, covering many aspects of the context.\n" + "5. Do not give away too much information in the questions. For example, ask 'Who is X' instead of 'Who is X that did Y' when Y is clear from the context.\n" + "6. Ignore the text formatting of the context, e.g., bold, italic, underline, etc.\n" + "7. Ignore typos, spacing and grammatical errors in the context.\n\n" + "Rules to follow when generate the answers:\n" + "1. The answers must use the information provided in the context.\n" + "2. Do not just copy words from the context. Answer the question in your own words.\n" + "3. The answers should be detailed and comprehensive. Please include additional specific details from the context when possible.\n\n" + "Response with {n_qa_pairs} question-answer pairs.\n" + "Always use proper grammar and punctuation.\n" + "Try to use different question forms and styles.\n" + "Use simple words and make sure that the answers are clear and comprehensive.\n\n" + "The question-answer pairs should be in the following format:\n" + "Question 1: {{question_1}}\n" + "Answer 1: {{answer_1}}\n" + "Question 2: {{question_2}}\n" + "Answer 2: {{answer_2}}\n" + "..." +) + + +def get_prompt(context, n_qa_pairs): + prompt = PROMPT_TEMPLATE.format(context=context, n_qa_pairs=n_qa_pairs) + return prompt + + +def postprocess_qa_pairs(res_txt: str): + """ + Postprocesses the QA pairs from the response text. + + Args: + res_txt: The response text. + n_qa_pairs: The number of QA pairs. + + Returns: + A tuple of two lists, the first containing the questions and the second containing the answers. + """ + # capture everything after each "Question {number}:" until "Answer" + q_pattern = r"Question \d+:(.*?)(?=Answer|$)" # thanks chatgpt + questions = re.findall(q_pattern, res_txt, flags=re.S) + + a_pattern = r"Answer \d+:(.*?)(?=Question|$)" # thanks chatgpt + answers = re.findall(a_pattern, res_txt, flags=re.S) + + if len(questions) != len(answers): + print(f"Warning---number of questions and answers do not match") + print(f"Number of questions: {len(questions)}") + print(f"Number of answers: {len(answers)}") + + out_q = [] + out_a = [] + if (len(questions) > 0) and (len(answers) > 0): + for i in range(min(len(questions), len(answers))): + out_q.append(questions[i].strip()) + out_a.append(answers[i].strip()) + + return out_q, out_a + + +def length_filter(samples): + return [len(text) < 10_000 for text in samples["text"]] + + +if __name__ == "__main__": + # api_key = os.environ.get("vllm_api_key") + # vllm_port = os.environ.get("vllm_port") + # vllm_ip = os.environ.get("vllm_ip") # "172.16.0.62" + vllm_model = os.environ.get("vllm_model") # "google/gemma-2-27b-it" + print(f"Using model: {vllm_model}") + # print(f"Using API key: {api_key}") + # print(f"Using VLLM IP: {vllm_ip}") + # print(f"Using VLLM port: {vllm_port}") + # client = OpenAI( + # base_url=f"http://{vllm_ip}:{vllm_port}/v1", + # api_key=api_key, + # ) + mm_kwargs = {"image":0} # if "gemma-3" in vllm_model else {} + llm_kwargs = dict( + model=vllm_model, + dtype="bfloat16", + # tokenizer_mode="mistral", + # config_format="mistral", + # load_format="mistral", + enable_prefix_caching=True, + max_model_len=11264,#2**14, + limit_mm_per_prompt=mm_kwargs, + # enable_chunked_prefill=True, + ) + + if vllm_model=="mistralai/Mistral-Small-3.1-24B-Instruct-2503": + llm_kwargs["tokenizer_mode"] = "mistral" + llm_kwargs["config_format"] = "mistral" + llm_kwargs["load_format"] = "mistral" + SYSTEM_TEMPLATE = ( + "You are Mistral Small 3.1, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\n" + "You power an AI assistant called Le Chat.\n" + "Your knowledge base was last updated on 2023-10-01.\n\n\n" + ) + SYSTEM_TEMPLATE + llm = LLM(**llm_kwargs) + tokenizer = llm.get_tokenizer() + shard_pattern = sys.argv[1] + n_qa_pairs = int(sys.argv[2]) + + paths = glob(f"./data/raw_datasets/fineweb_edu/sample/100BT/{shard_pattern}.parquet") + + for path in paths: + ds = load_dataset( + "parquet", + data_files=path, + split="train", + streaming=True, + ) + ds = ds.filter(length_filter, batched=True) + + ctxs = [sample["text"] for sample in iter(ds)] + messages = [ + [ + {"role": "system", "content": SYSTEM_TEMPLATE}, + {"role": "user", "content": get_prompt(ctx, n_qa_pairs)}, + ] + for ctx in ctxs + ] + + print(f"Generating from {len(messages)} contexts") + completions = llm.chat( + messages, + sampling_params=SamplingParams( + max_tokens=2048, + temperature=0.1 if vllm_model=="mistralai/Mistral-Small-3.1-24B-Instruct-2503" else 0.7, + # frequency_penalty=0.1, + ), + ) + samples = [] + for ctx, completion in zip(ctxs, completions): + questions, answers = postprocess_qa_pairs(completion.outputs[0].text) + for q, a in zip(questions, answers): + samples.append( + { + "context": ctx, + "prompt": q, + "response": a, + } + ) + + print(f"Generated {len(samples)} samples") + random.shuffle(samples) + df = pd.DataFrame(samples) + ds = Dataset.from_pandas(df) + val_ds = ds.take(10) + ds = ds.skip(10) + + shard_name = path.split("/")[-1].split(".")[0] + ds.to_parquet(f"data/raw_datasets/fw_qa_3/{shard_name}.parquet") + val_ds.to_parquet(f"data/raw_datasets/fw_qa_3/{shard_name}_val.parquet") + print(f"Saved to data/raw_datasets/fw_qa_3/{shard_name}.parquet") + print(f"Saved to data/raw_datasets/fw_qa_3/{shard_name}_val.parquet") + + # debug + # ds = load_dataset( + # "parquet", + # data_files=paths[0], + # split="train[:10]", + # ) + # ctxs = [sample["text"] for sample in iter(ds)] + # messages = [ + # [ + # {"role": "system", "content": SYSTEM_TEMPLATE}, + # {"role": "user", "content": get_prompt(ctx, n_qa_pairs)}, + # ] + # for ctx in ctxs + # ] + # + # print(f"Generating from {len(messages)} contexts") + # completions = llm.chat( + # messages, + # sampling_params=SamplingParams( + # max_tokens=2048, + # temperature=0.1 if vllm_model=="mistralai/Mistral-Small-3.1-24B-Instruct-2503" else 0.7, + # # frequency_penalty=0.1, + # ), + # ) + # samples = [] + # for ctx, completion in zip(ctxs, completions): + # questions, answers = postprocess_qa_pairs(completion.outputs[0].text) + # for q, a in zip(questions, answers): + # samples.append( + # { + # "context": ctx, + # "prompt": q, + # "response": a, + # } + # ) + # for sample in samples: + # print(f"{sample['context']=}") + # print(f"{sample['prompt']=}") + # print(f"{sample['response']=}") + # print() diff --git a/generate_fw_qa_vllm.py b/generate_fw_qa_vllm.py index 919bfe2..3c5349a 100644 --- a/generate_fw_qa_vllm.py +++ b/generate_fw_qa_vllm.py @@ -13,34 +13,44 @@ from openai import OpenAI SYSTEM_TEMPLATE = ( "You are a creative and helpful assistant.\n" - "You are given a context and you need to generate questions and corresponding answers from the given context.\n" + "You are tasked with generating questions for reading comprehension tests.\n" + "You will be given a context and you need to generate questions and corresponding answers from the given context.\n" "The questions should be highly specific to the information provided in the context, not general questions that suits any context.\n" - "Do not halucinate and make up information." + "**DO NOT** halucinate or make up information." ) # based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811) PROMPT_TEMPLATE = ( + "### Instructions ###\n" "Generate questions and corresponding answers from the given context. The questions should be highly specific to the " - "information provided in the context, not general questions that suits any context.\n\n" + "information provided in the context, not general questions that suits any context.\n\n\n" + "### Context ###\n" + "{context}" + "\n\n\n" + "### Additional Instructions ###\n" "Rules to follow when generate the questions:\n" - "1. The questions must be fully answerable from information present in given context.\n" + "1. The questions must be specific to the given context and fully answerable from information present in given context.\n" "2. Make sure the questions are clear and unambiguous.\n" - "3. Phrases like 'based on the provided context', 'according to the context', etc, are not allowed to appear in " - "the questions.\n\n" + "3. Phrases like 'based on the provided context', 'according to the context', etc, are **NOT ALLOWED** to appear in " + "the questions.\n" + "4. The questions should not overlap. They should be diverse, covering many aspects of the context.\n" + "5. Do not give away too much information in the questions. For example, ask 'Who is X' instead of 'Who is X that did Y' when Y is clear from the context.\n" + "6. Ignore the text formatting of the context, e.g., bold, italic, underline, etc.\n" + "7. Ignore typos, spacing and grammatical errors in the context.\n\n" "Rules to follow when generate the answers:\n" "1. The answers must use the information provided in the context.\n" "2. Do not just copy words from the context. Answer the question in your own words.\n" - "3. The answers should be detailed and comprehensive.\n\n" - "Response with {n_qa_pairs} question-answer pairs. Use simple words and please be clear.\n" + "3. The answers should be detailed and comprehensive. Please include additional specific details from the context when possible.\n\n" + "Response with {n_qa_pairs} question-answer pairs.\n" + "Always use proper grammar and punctuation.\n" + "Try to use different question forms and styles.\n" + "Use simple words and make sure that the answers are clear and comprehensive.\n\n" "The question-answer pairs should be in the following format:\n" "Question 1: {{question_1}}\n" "Answer 1: {{answer_1}}\n" "Question 2: {{question_2}}\n" "Answer 2: {{answer_2}}\n" "..." - "\n\n" - "### Context ###\n" - "{context}" ) @@ -82,6 +92,10 @@ def postprocess_qa_pairs(res_txt: str): return out_q, out_a +def length_filter(samples): + return [len(text) < 10_000 for text in samples["text"]] + + if __name__ == "__main__": # api_key = os.environ.get("vllm_api_key") # vllm_port = os.environ.get("vllm_port") @@ -95,79 +109,123 @@ if __name__ == "__main__": # base_url=f"http://{vllm_ip}:{vllm_port}/v1", # api_key=api_key, # ) - llm = LLM( + mm_kwargs = {"image":0} # if "gemma-3" in vllm_model else {} + llm_kwargs = dict( model=vllm_model, - tokenizer_mode="mistral", - config_format="mistral", - load_format="mistral", + # tokenizer_mode="mistral", + # config_format="mistral", + # load_format="mistral", enable_prefix_caching=True, + max_model_len=12288,#2**14, + limit_mm_per_prompt=mm_kwargs, # enable_chunked_prefill=True, ) + + if vllm_model=="mistralai/Mistral-Small-3.1-24B-Instruct-2503": + llm_kwargs["tokenizer_mode"] = "mistral" + llm_kwargs["config_format"] = "mistral" + llm_kwargs["load_format"] = "mistral" + SYSTEM_TEMPLATE = ( + "You are Mistral Small 3.1, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\n" + "You power an AI assistant called Le Chat.\n" + "Your knowledge base was last updated on 2023-10-01.\n\n\n" + ) + SYSTEM_TEMPLATE + llm = LLM(**llm_kwargs) tokenizer = llm.get_tokenizer() shard_pattern = sys.argv[1] n_qa_pairs = int(sys.argv[2]) - for path in glob(f"./data/raw_datasets/fineweb_sharded/{shard_pattern}.parquet"): - ds = load_dataset( - "parquet", - data_files=path, - split="train", - streaming=True, - ) + paths = glob(f"./data/raw_datasets/fineweb/sample/100BT/{shard_pattern}.parquet") - ctxs = [sample["text"] for sample in iter(ds)] - messages = [ - [ - {"role": "system", "content": SYSTEM_TEMPLATE}, - {"role": "user", "content": get_prompt(ctx, n_qa_pairs)}, - ] - for ctx in ctxs - ] - - print(f"Generating from {len(messages)} contexts") - completions = llm.chat( - messages, - sampling_params=SamplingParams( - max_tokens=2048, - temperature=1.0, - frequency_penalty=0.2, - ), - ) - samples = [] - for ctx, completion in zip(ctxs, completions): - questions, answers = postprocess_qa_pairs(completion.outputs[0].text) - for q, a in zip(questions, answers): - samples.append( - { - "context": ctx, - "prompt": q, - "response": a, - } - ) - - print(f"Generated {len(samples)} samples") - random.shuffle(samples) - df = pd.DataFrame(samples) - ds = Dataset.from_pandas(df) - val_ds = ds.take(10) - ds = ds.skip(10) - - shard_name = path.split("/")[-1].split(".")[0] - ds.to_parquet(f"data/raw_datasets/fw_qa_xl/{shard_name}.parquet") - val_ds.to_parquet(f"data/raw_datasets/fw_qa_xl/{shard_name}_val.parquet") - print(f"Saved to data/raw_datasets/fw_qa_xl/{shard_name}.parquet") - print(f"Saved to data/raw_datasets/fw_qa_xl/{shard_name}_val.parquet") - - # for i, sample in tqdm(enumerate(iter(ds))): - # completion = client.chat.completions.create( - # model=vllm_model, - # messages=[ - # {"role": "system", "content": SYSTEM_TEMPLATE}, - # {"role": "user", "content": get_prompt(sample["text"], n_qa_pairs)}, - # ], - # extra_body={"temperature": 1.0, "frequency_penalty": 0.2}, + # for path in paths: + # ds = load_dataset( + # "parquet", + # data_files=path, + # split="train", + # streaming=True, # ) + # ds = ds.filter(length_filter, batched=True) + # + # ctxs = [sample["text"] for sample in iter(ds)] + # messages = [ + # [ + # {"role": "system", "content": SYSTEM_TEMPLATE}, + # {"role": "user", "content": get_prompt(ctx, n_qa_pairs)}, + # ] + # for ctx in ctxs + # ] + # + # print(f"Generating from {len(messages)} contexts") + # completions = llm.chat( + # messages, + # sampling_params=SamplingParams( + # max_tokens=2048, + # temperature=0.1 if vllm_model=="mistralai/Mistral-Small-3.1-24B-Instruct-2503" else 0.7, + # # frequency_penalty=0.1, + # ), + # ) + # samples = [] + # for ctx, completion in zip(ctxs, completions): + # questions, answers = postprocess_qa_pairs(completion.outputs[0].text) + # for q, a in zip(questions, answers): + # samples.append( + # { + # "context": ctx, + # "prompt": q, + # "response": a, + # } + # ) + # + # print(f"Generated {len(samples)} samples") + # random.shuffle(samples) + # df = pd.DataFrame(samples) + # ds = Dataset.from_pandas(df) + # val_ds = ds.take(10) + # ds = ds.skip(10) + # + # shard_name = path.split("/")[-1].split(".")[0] + # ds.to_parquet(f"data/raw_datasets/fw_qa_3/{shard_name}.parquet") + # val_ds.to_parquet(f"data/raw_datasets/fw_qa_3/{shard_name}_val.parquet") + # print(f"Saved to data/raw_datasets/fw_qa_3/{shard_name}.parquet") + # print(f"Saved to data/raw_datasets/fw_qa_3/{shard_name}_val.parquet") - # print(completion.choices[0].message) - # if i >= 10: - # break + # debug + ds = load_dataset( + "parquet", + data_files=paths[0], + split="train[:10]", + ) + ctxs = [sample["text"] for sample in iter(ds)] + messages = [ + [ + {"role": "system", "content": SYSTEM_TEMPLATE}, + {"role": "user", "content": get_prompt(ctx, n_qa_pairs)}, + ] + for ctx in ctxs + ] + + print(f"Generating from {len(messages)} contexts") + completions = llm.chat( + messages, + sampling_params=SamplingParams( + max_tokens=2048, + temperature=0.1 if vllm_model=="mistralai/Mistral-Small-3.1-24B-Instruct-2503" else 0.7, + # frequency_penalty=0.1, + ), + ) + samples = [] + for ctx, completion in zip(ctxs, completions): + questions, answers = postprocess_qa_pairs(completion.outputs[0].text) + for q, a in zip(questions, answers): + samples.append( + { + "context": ctx, + "prompt": q, + "response": a, + } + ) + for sample in samples: + print(f"{sample['context']=}") + print(f"{sample['prompt']=}") + print(f"{sample['response']=}") + print() diff --git a/intx_sft.py b/intx_sft.py index 39332d7..d551b7a 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -573,5 +573,5 @@ if __name__ == "__main__": if os.getenv("DEBUG", False): disable_caching() # randomly sleep to avoid run_name collision - time.sleep(random.random() * 13) + # time.sleep(random.random() * 13) main() diff --git a/process_fineweb.py b/process_fineweb.py index 58f03b8..b8a6ac2 100644 --- a/process_fineweb.py +++ b/process_fineweb.py @@ -17,23 +17,23 @@ if __name__ == "__main__": "HuggingFaceFW/fineweb", repo_type="dataset", local_dir=fw_dir, - allow_patterns="sample/10BT/*", + allow_patterns="sample/100BT/*", ) - # https://github.com/huggingface/datasets/issues/7047#issuecomment-2233163406 - num_shards_per_file = 16 - sharded_fw_dir = "./data/raw_datasets/fineweb_sharded/" - output_path_template = f"{sharded_fw_dir}" + "/{i:02d}_{idx:05d}.parquet" + # # https://github.com/huggingface/datasets/issues/7047#issuecomment-2233163406 + # num_shards_per_file = 16 + # sharded_fw_dir = "./data/raw_datasets/fineweb_sharded/" + # output_path_template = f"{sharded_fw_dir}" + "/{i:02d}_{idx:05d}.parquet" - for i, f in enumerate(sorted(glob(f"{fw_dir}/sample/10BT/*.parquet"))): - # ~1M rows ~= 2GB mem required per file - ds = Dataset.from_parquet(f) - ds = ds.filter(remove_too_long, batched=True) - print(f"Filtered ds size: {len(ds)}") - ds = ds.shuffle(seed=42 + i) - # take one shard (from 16 shards) per file - # idx = random.sample(range(num_shards_per_file), 1)[0] - # shard = ds.shard(index=idx, num_shards=num_shards_per_file, contiguous=False) - # shard.to_parquet(output_path_template.format(index=i)) - for idx in range(num_shards_per_file): - shard = ds.shard(index=idx, num_shards=num_shards_per_file, contiguous=True) - shard.to_parquet(output_path_template.format(i=i, idx=idx)) + # for i, f in enumerate(sorted(glob(f"{fw_dir}/sample/100BT/*.parquet"))): + # # ~1M rows ~= 2GB mem required per file + # ds = Dataset.from_parquet(f) + # ds = ds.filter(remove_too_long, batched=True) + # print(f"Filtered ds size: {len(ds)}") + # ds = ds.shuffle(seed=42 + i) + # # take one shard (from 16 shards) per file + # # idx = random.sample(range(num_shards_per_file), 1)[0] + # # shard = ds.shard(index=idx, num_shards=num_shards_per_file, contiguous=False) + # # shard.to_parquet(output_path_template.format(index=i)) + # for idx in range(num_shards_per_file): + # shard = ds.shard(index=idx, num_shards=num_shards_per_file, contiguous=True) + # shard.to_parquet(output_path_template.format(i=i, idx=idx)) diff --git a/requirements.txt b/requirements.txt index 472156f..cf54865 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,7 +2,8 @@ einops jaxtyping liger-kernel rouge-score -transformers==4.46.3 +transformers==4.51.1 +# vllm==0.8.3 deepspeed==0.15.4 # curretly 0.16.2 accelerate==1.2.1 # visualization diff --git a/scripts/repeat_text/gemma_llama_xl.sh b/scripts/repeat_text/gemma_llama_xl.sh new file mode 100644 index 0000000..cb52901 --- /dev/null +++ b/scripts/repeat_text/gemma_llama_xl.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=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=True \ +--use_sequence_packing=True --ctx_encoder_model_name_or_path=meta-llama/Llama-3.2-3B-Instruct diff --git a/scripts/scale_modules_fixed_latents/gemma_llama_xl.sh b/scripts/scale_modules_fixed_latents/gemma_llama_xl.sh new file mode 100644 index 0000000..cb6c811 --- /dev/null +++ b/scripts/scale_modules_fixed_latents/gemma_llama_xl.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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 29561 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_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--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 \ +--use_sequence_packing=True diff --git a/scripts/scale_modules_fixed_latents/gemma_llama_xl_down_ln.sh b/scripts/scale_modules_fixed_latents/gemma_llama_xl_down_ln.sh new file mode 100644 index 0000000..a75e91c --- /dev/null +++ b/scripts/scale_modules_fixed_latents/gemma_llama_xl_down_ln.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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 29566 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 --extra_modules=input_layernorm,post_attention_layernorm \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=2 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--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 \ +--use_sequence_packing=True diff --git a/scripts/scale_modules_fixed_latents/gemma_llama_xl_ln_only.sh b/scripts/scale_modules_fixed_latents/gemma_llama_xl_ln_only.sh new file mode 100644 index 0000000..8506a41 --- /dev/null +++ b/scripts/scale_modules_fixed_latents/gemma_llama_xl_ln_only.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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=4 --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=4 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=None --extra_modules=input_layernorm,post_attention_layernorm \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=4 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--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 \ +--use_sequence_packing=True diff --git a/scripts/scale_modules_fixed_latents/gemma_llama_xl_up_down.sh b/scripts/scale_modules_fixed_latents/gemma_llama_xl_up_down.sh new file mode 100644 index 0000000..b78380c --- /dev/null +++ b/scripts/scale_modules_fixed_latents/gemma_llama_xl_up_down.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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,up_proj \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=4 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--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 \ +--use_sequence_packing=True diff --git a/batch_gemma_llama_new.sh b/scripts/scale_modules_fixed_latents/gemma_llama_xl_up_down_ln.sh similarity index 51% rename from batch_gemma_llama_new.sh rename to scripts/scale_modules_fixed_latents/gemma_llama_xl_up_down_ln.sh index d6db711..cf90799 100644 --- a/batch_gemma_llama_new.sh +++ b/scripts/scale_modules_fixed_latents/gemma_llama_xl_up_down_ln.sh @@ -2,7 +2,7 @@ #SBATCH --job-name=gemma_llama_instruct #SBATCH --partition=a3 #SBATCH --nodes=1 -#SBATCH --gpus=8 +#SBATCH --gpus=4 #SBATCH --output=outputs/%x-%j.out #SBATCH --error=outputs/%x-%j.out @@ -15,17 +15,17 @@ # 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/llm-modulator +conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora # eval "$@" -accelerate launch --num_processes=8 --gradient_accumulation_steps=1 --gradient_clipping=1.0 \ ---gpu_ids all --main_process_port 29554 intx_sft.py configs/pretrain_all.yaml \ +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=10.1 --per_device_train_batch_size=32 \ ---gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ ---target_modules=down_proj --num_blocks=1 --num_self_attends_per_block=16 \ ---eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 \ ---neftune_noise_alpha=1 --use_light_weight_lora=True --light_weight_latent_size=512 \ +--gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj,up_proj --extra_modules=input_layernorm,post_attention_layernorm \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=2 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--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.1-8B-Instruct \ +--add_repeat_prompt=False \ --use_sequence_packing=True - diff --git a/scripts/scale_modules_no_lightweight/gemma_llama_2.sh b/scripts/scale_modules_no_lightweight/gemma_llama_2.sh new file mode 100644 index 0000000..3b5a04b --- /dev/null +++ b/scripts/scale_modules_no_lightweight/gemma_llama_2.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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_2.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1.0 --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_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=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 --ctx_encoder_model_name_or_path=meta-llama/Llama-3.2-3B-Instruct diff --git a/scripts/scale_modules_no_lightweight/gemma_llama_xl.sh b/scripts/scale_modules_no_lightweight/gemma_llama_xl.sh new file mode 100644 index 0000000..bb55381 --- /dev/null +++ b/scripts/scale_modules_no_lightweight/gemma_llama_xl.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=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 --ctx_encoder_model_name_or_path=meta-llama/Llama-3.2-3B-Instruct diff --git a/scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen.sh b/scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen.sh new file mode 100644 index 0000000..4605e1c --- /dev/null +++ b/scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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=16 \ +--gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=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 diff --git a/scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen_latent_fac_1.sh b/scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen_latent_fac_1.sh new file mode 100644 index 0000000..70de581 --- /dev/null +++ b/scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen_latent_fac_1.sh @@ -0,0 +1,32 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=4 \ +--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 diff --git a/scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen_latent_fac_1_llama.sh b/scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen_latent_fac_1_llama.sh new file mode 100644 index 0000000..7dbb986 --- /dev/null +++ b/scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen_latent_fac_1_llama.sh @@ -0,0 +1,33 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=2 \ +--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 \ +--ctx_encoder_model_name_or_path=meta-llama/Llama-3.2-3B-Instruct + diff --git a/scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen_latent_fac_1_r_16.sh b/scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen_latent_fac_1_r_16.sh new file mode 100644 index 0000000..4421159 --- /dev/null +++ b/scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen_latent_fac_1_r_16.sh @@ -0,0 +1,32 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=1 \ +--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 diff --git a/scripts/scale_modules_no_lightweight/gemma_llama_xl_r_16.sh b/scripts/scale_modules_no_lightweight/gemma_llama_xl_r_16.sh new file mode 100644 index 0000000..2e8ae06 --- /dev/null +++ b/scripts/scale_modules_no_lightweight/gemma_llama_xl_r_16.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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=16 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29564 intx_sft.py configs/pretrain_all_xl_and_sum.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=32 \ +--gradient_accumulation_steps=16 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=16 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_r=16 --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 diff --git a/scripts/scale_modules_no_lightweight/gemma_llama_xl_r_16_ln.sh b/scripts/scale_modules_no_lightweight/gemma_llama_xl_r_16_ln.sh new file mode 100644 index 0000000..fda30d8 --- /dev/null +++ b/scripts/scale_modules_no_lightweight/gemma_llama_xl_r_16_ln.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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=16 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29563 intx_sft.py configs/pretrain_all_xl_and_sum.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=32 \ +--gradient_accumulation_steps=16 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj --extra_modules=input_layernorm,post_attention_layernorm \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=16 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_r=16 --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 diff --git a/scripts/scale_modules_no_lightweight/gemma_llama_xl_up_down.sh b/scripts/scale_modules_no_lightweight/gemma_llama_xl_up_down.sh new file mode 100644 index 0000000..fa3ce26 --- /dev/null +++ b/scripts/scale_modules_no_lightweight/gemma_llama_xl_up_down.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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=16 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29563 intx_sft.py configs/pretrain_all_xl_and_sum.yaml \ +--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5.1 --per_device_train_batch_size=32 \ +--gradient_accumulation_steps=16 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=up_proj,down_proj \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_r=8 --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 diff --git a/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl.sh b/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl.sh new file mode 100644 index 0000000..6614d86 --- /dev/null +++ b/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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=4 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29571 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=4 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=16 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=3e-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 diff --git a/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_down_ln.sh b/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_down_ln.sh new file mode 100644 index 0000000..265f185 --- /dev/null +++ b/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_down_ln.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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=4 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29565 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=4 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj --extra_modules=input_layernorm,post_attention_layernorm \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--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 \ +--use_sequence_packing=True diff --git a/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_ln_only.sh b/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_ln_only.sh new file mode 100644 index 0000000..785207c --- /dev/null +++ b/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_ln_only.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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=4 --gradient_clipping=1.0 \ +--gpu_ids all --main_process_port 29564 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=4 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=None --extra_modules=input_layernorm,post_attention_layernorm \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--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 \ +--use_sequence_packing=True diff --git a/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_up_down.sh b/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_up_down.sh new file mode 100644 index 0000000..dbe9e5e --- /dev/null +++ b/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_up_down.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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=4 --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=4 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj,up_proj \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--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 \ +--use_sequence_packing=True diff --git a/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_up_down_ln.sh b/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_up_down_ln.sh new file mode 100644 index 0000000..ba763f6 --- /dev/null +++ b/scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_up_down_ln.sh @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=gemma_llama_instruct +#SBATCH --partition=a3 +#SBATCH --nodes=1 +#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=4 --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=4 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \ +--target_modules=down_proj,up_proj --extra_modules=input_layernorm,post_attention_layernorm \ +--num_blocks=1 --num_self_attends_per_block=8 --num_latent_factor=8 \ +--eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --lora_dropout=0.0 \ +--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 \ +--use_sequence_packing=True diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index 98348d3..9960b30 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -374,6 +374,10 @@ class HypernetArguments: default=None, metadata={"help": "Extra modules to train."}, ) + per_rank_gen: bool = field( + default=False, + metadata={"help": "Whether to use per-rank generation."}, + ) @dataclass diff --git a/src/ctx_to_lora/data_utils.py b/src/ctx_to_lora/data_utils.py index 64692b0..3ba1da9 100644 --- a/src/ctx_to_lora/data_utils.py +++ b/src/ctx_to_lora/data_utils.py @@ -86,6 +86,18 @@ DS_KWARGS = { split="train", ), ), + "fw_qa_2": dict( + train=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_2/*[!val].parquet"), + split="train", + ), + validation=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_2/*val.parquet"), + split="train", + ), + ), "ctx_qa": dict( train=dict( path="parquet", @@ -152,6 +164,20 @@ DS_KWARGS = { split="test", ), ), + # "wikitext-2": dict( + # train=dict( + # path="EleutherAI/wikitext_document_level", + # name="wikitext-2-raw-v1", + # split="train", + # ), + # ), + # "wikitext-103": dict( + # train=dict( + # path="EleutherAI/wikitext_document_level", + # name="wikitext-103-raw-v1", + # split="train", + # ), + # ), "gsm8k": dict( train=dict( path="openai/gsm8k", @@ -224,11 +250,11 @@ def add_repeat_prompt_fn(samples): ): # Only process if the context is not already in the set if ctx in unique_contexts: + continue + # remove too long contexts + if len(ctx) > 1000: continue unique_contexts.add(ctx) - # remove too long contexts - if len(ctx) > 3000: - continue ctxs.append(ctx) responses.append(ctx) prompts.append("Repeat the text above.") @@ -341,9 +367,9 @@ def get_tokenized_dataset( ds = ds.filter(filter_long_samples, batched=True, num_proc=16) if split == "train": if add_negative_prompt: - ds = ds.map(add_negative_prompt_fn, batched=True, batch_size=None) + ds = ds.map(add_negative_prompt_fn, batched=True) if add_repeat_prompt and "context_numbers" not in ds_name: - ds = ds.map(add_repeat_prompt_fn, batched=True, batch_size=None) + ds = ds.map(add_repeat_prompt_fn, batched=True) tokenized_ds = construct_and_tokenize_ctx_qa( tokenizer, tokenizer_kwargs, @@ -612,6 +638,15 @@ def get_preprocessing_fn(ds_name: str) -> Callable[[dict[str, Any]], dict[str, A "response": sample["summary"], } + elif "wikitext" in ds_name: + + def f(sample): + return { + "context": sample["page"], + "prompt": "PLAECHOLDER", + "response": "PLAECHOLDER", + } + elif ds_name == "openmathintx-2": def f(sample): diff --git a/src/ctx_to_lora/eval.py b/src/ctx_to_lora/eval.py index aef2374..fa72d7f 100644 --- a/src/ctx_to_lora/eval.py +++ b/src/ctx_to_lora/eval.py @@ -574,14 +574,14 @@ if __name__ == "__main__": else: args.test_ds_names = cli_args.datasets - evaluate( - cli_args.checkpoint_path, - cli_args.model_name_or_path, - cli_args.eval_batch_size, - args, - split=cli_args.split, - generative=False, - ) + # evaluate( + # cli_args.checkpoint_path, + # cli_args.model_name_or_path, + # cli_args.eval_batch_size, + # args, + # split=cli_args.split, + # generative=False, + # ) evaluate( cli_args.checkpoint_path, cli_args.model_name_or_path, diff --git a/src/ctx_to_lora/modeling_utils.py b/src/ctx_to_lora/modeling_utils.py index 04d7d5a..06389f4 100644 --- a/src/ctx_to_lora/modeling_utils.py +++ b/src/ctx_to_lora/modeling_utils.py @@ -76,6 +76,7 @@ class AggregatorConfig: feature_size: int num_layers: int num_modules: int + num_extra_modules: int output_size: int # perceiver @@ -85,6 +86,8 @@ class AggregatorConfig: self_attention_widening_factor: int = 4 cross_attention_widening_factor: int = 1 num_latent_factor: int = 8 + lora_r: int = 8 + per_rank_gen: bool = False def get_aggregator_config( @@ -92,6 +95,9 @@ def get_aggregator_config( ctx_encoder_model_config: PretrainedConfig, output_size: int, num_modules: int, + num_extra_modules: int, + lora_r: int, + per_rank_gen: bool, aggregator_args: AggregatorArguments, ): return AggregatorConfig( @@ -99,6 +105,9 @@ def get_aggregator_config( output_size=output_size, num_layers=get_num_layers(model), num_modules=num_modules, + num_extra_modules=num_extra_modules, + lora_r=lora_r, + per_rank_gen=per_rank_gen, **vars(aggregator_args), ) @@ -108,6 +117,7 @@ class HypernetConfig: latent_size: int use_light_weight_lora: bool light_weight_latent_size: int + per_rank_gen: bool dropout_rate: float lora_config: LoraConfig @@ -131,7 +141,7 @@ def get_hypernet_config( if lora_config is not None: lora_config = lora_config["default"] num_modules += len(lora_config.target_modules) - num_modules += len(hypernet_args.extra_modules or []) + num_extra_modules = len(hypernet_args.extra_modules or []) indices = torch.arange(get_num_layers(model), device=model.device) return HypernetConfig( **vars(hypernet_args), @@ -145,6 +155,9 @@ def get_hypernet_config( ctx_encoder_model_config, hypernet_args.latent_size, num_modules, + num_extra_modules, + lora_config.r, + hypernet_args.per_rank_gen, aggregator_args, ), ) @@ -159,6 +172,9 @@ class Perceiver(nn.Module): output_size, num_layers, num_modules, + num_extra_modules, + per_rank_gen, + lora_r, num_latent_factor, *args, **kwargs, @@ -166,6 +182,9 @@ class Perceiver(nn.Module): super().__init__() self.num_layers = num_layers self.num_modules = num_modules + self.num_extra_modules = num_extra_modules + self.per_rank_gen = per_rank_gen + self.r = lora_r if self.per_rank_gen else 1 # self.config = PerceiverConfig( # d_model=feature_size, # + num_bands # num_latents=num_layers * num_modules * num_latent_factor, @@ -193,7 +212,8 @@ class Perceiver(nn.Module): input_size=feature_size, # the first layer is xattn resampler_depth=kwargs["num_self_attends_per_block"] + 1, - resampler_n_latents=num_layers * num_modules * num_latent_factor, + resampler_n_latents=(num_layers * (num_modules * self.r + num_extra_modules)) + * num_latent_factor, intermediate_size_factor=4, hidden_size=output_size, attn_implementation="flash_attention_2", @@ -201,7 +221,7 @@ class Perceiver(nn.Module): self.decoder_config = Idefics2PerceiverConfig( input_size=output_size, resampler_depth=1, - resampler_n_latents=num_layers * num_modules, + resampler_n_latents=num_layers * (num_modules * self.r + num_extra_modules), hidden_size=output_size, attn_implementation="flash_attention_2", ) @@ -214,13 +234,43 @@ class Perceiver(nn.Module): ctx_position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None, ): x = self.perceiver(ctx_features, ctx_attn_mask, ctx_position_ids) - x = rearrange( + lora_x, extra_x = unpack( x, - "bs (n_layers n_modules) d -> bs n_layers n_modules d", + [ + [self.num_layers * self.num_modules * self.r], + [self.num_layers * self.num_extra_modules], + ], + "bs * feature_dim", + ) + lora_x = rearrange( + lora_x, + "bs (n_layers n_modules r) d -> bs n_layers n_modules r d", n_modules=self.num_modules, n_layers=self.num_layers, + r=self.r, ) - return x + if not self.per_rank_gen: + lora_x = lora_x.squeeze(3) + + extra_x = rearrange( + extra_x, + "bs (n_layers n_extra_modules) d -> bs n_layers n_extra_modules d", + n_extra_modules=self.num_extra_modules, + n_layers=self.num_layers, + ) + + # x = rearrange( + # x, + # "bs (n_layers n_modules) d -> bs n_layers n_modules d", + # n_modules=self.num_modules, + # n_layers=self.num_layers, + # ) + # lora_emb, extra_emb = unpack( + # emb, + # [[self.num_modules], [self.num_extra_modules]], + # "bs n_layers * feature_dim", + # ) + return lora_x, extra_x class Mixer(nn.Module): @@ -485,12 +535,24 @@ class HyperLoRA(nn.Module): self.d_in, self.d_out = self.config.feature_sizes if self.target_modules: + # self.layers = nn.Sequential( + # *[ + # 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), + # ) + # for _ in range(4) + # ] + # ) self.layers = 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), ) + # self.layers = nn.Identity() if self.config.use_light_weight_lora: # light-weight lora projection (per module) @@ -535,36 +597,73 @@ class HyperLoRA(nn.Module): # has something to do with the bias shape (n_modules r d_lora) # when n_modules == 1, adamw_torch_fused complains about device/layout # but when n_modules > 1, it works fine - if n_modules == 1: - self.head = Mix( - "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", - weight_shape="d_latent r d_lora", - # bias_shape=None, # no bias - bias_shape="r d_lora", - d_latent=self.config.latent_size, - r=self.config.lora_config.r, - d_lora=d_lora, - ) + if self.config.per_rank_gen: + if n_modules == 1: + self.head = Mix( + "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", + weight_shape="d_latent d_lora", + # bias_shape=None, # no bias + bias_shape="d_lora", + # n_layers=len(self.layer_indices), + d_latent=self.config.latent_size, + r=self.config.lora_config.r, + d_lora=d_lora, + ) + else: + self.head = Mix( + "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", + weight_shape="n_modules d_latent d_lora", + # bias_shape=None, # no bias + bias_shape="n_modules d_lora", + # n_layers=len(self.layer_indices), + n_modules=n_modules, + d_latent=self.config.latent_size, + r=self.config.lora_config.r, + d_lora=d_lora, + ) else: - # each module processes d -> r d_out independently - self.head = Mix( - "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", - weight_shape="n_modules d_latent r d_lora", - # bias_shape=None, # no bias - bias_shape="n_modules r d_lora", - n_modules=len(self.target_modules), - d_latent=self.config.latent_size, - r=self.config.lora_config.r, - d_lora=d_lora, - ) + if n_modules == 1: + self.head = Mix( + "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", + weight_shape="d_latent r d_lora", + # bias_shape=None, # no bias + bias_shape="r d_lora", + d_latent=self.config.latent_size, + r=self.config.lora_config.r, + d_lora=d_lora, + ) + else: + # each module processes d -> r d_out independently + self.head = Mix( + "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", + weight_shape="n_modules d_latent r d_lora", + # bias_shape=None, # no bias + bias_shape="n_modules r d_lora", + n_modules=n_modules, + d_latent=self.config.latent_size, + r=self.config.lora_config.r, + d_lora=d_lora, + ) if self.extra_modules: + # self.extra_layers = nn.Sequential( + # *[ + # 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), + # ) + # for _ in range(4) + # ] + # ) self.extra_layers = 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), ) + # self.extra_layers = nn.Identity() # self.pre_extra_layers_norm = nn.LayerNorm(self.config.latent_size) # self.extra_norm = nn.LayerNorm(self.config.latent_size) @@ -649,12 +748,12 @@ class HyperLoRA(nn.Module): ): # [bs, n_layers, n_total_modules, feature_dim] - emb = self.aggregator(features, attn_mask, position_ids) - lora_emb, extra_emb = unpack( - emb, - [[self.num_modules], [self.num_extra_modules]], - "bs n_layers * feature_dim", - ) + lora_emb, extra_emb = self.aggregator(features, attn_mask, position_ids) + # lora_emb, extra_emb = unpack( + # emb, + # [[self.num_modules], [self.num_extra_modules]], + # "bs n_layers * feature_dim", + # ) # [bs, n_layers, n_modules, r, max_in_d_outim] flat_loras = None @@ -815,6 +914,9 @@ class ModulatedPretrainedModel(nn.Module): if self.hypernet.config.use_light_weight_lora: A = A[:, : self.hypernet.config.light_weight_latent_size] B = B[: self.hypernet.config.light_weight_latent_size] + if self.hypernet.config.per_rank_gen: + A = A[0:1] + B = B[:, 0:1] biases = [A, B.T] # bias-hyperinit