diff --git a/.gitignore b/.gitignore index 5ed6ddb..a5669eb 100644 --- a/.gitignore +++ b/.gitignore @@ -19,4 +19,5 @@ plots/ *.pth *.bin *.safetensors -*tfevents* \ No newline at end of file +*tfevents* +watcher_state.yaml diff --git a/batch.sh b/batch.sh new file mode 100755 index 0000000..9e1a41d --- /dev/null +++ b/batch.sh @@ -0,0 +1,29 @@ +#!/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 new file mode 100755 index 0000000..dcc8377 --- /dev/null +++ b/batch_gemma_2_2b.sh @@ -0,0 +1,30 @@ +#!/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_vision.sh b/batch_gemma_llama_vision.sh new file mode 100755 index 0000000..03ff5c8 --- /dev/null +++ b/batch_gemma_llama_vision.sh @@ -0,0 +1,30 @@ +#!/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 new file mode 100755 index 0000000..76ba443 --- /dev/null +++ b/batch_llama_2_7b.sh @@ -0,0 +1,29 @@ +#!/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_3_8b_1k_latent.sh b/batch_llama_3_8b_1k_latent.sh new file mode 100755 index 0000000..374af09 --- /dev/null +++ b/batch_llama_3_8b_1k_latent.sh @@ -0,0 +1,29 @@ +#!/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=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,up_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=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/generate_fw_qa.py b/generate_fw_qa.py index e709d13..9a2d2f3 100644 --- a/generate_fw_qa.py +++ b/generate_fw_qa.py @@ -93,30 +93,32 @@ def parse_args() -> Namespace: if __name__ == "__main__": # Load the datasets args = parse_args() - ds = load_dataset( - "parquet", - data_files=f"./data/raw_datasets/fineweb_sharded/{args.shard_pattern}.parquet", - split="train", - streaming=True, - ) - os.makedirs("openai_batches", exist_ok=True) - lines = [] - c = 0 - for i, sample in tqdm(enumerate(iter(ds))): - if len(lines) >= 20_000: - with open(f"openai_batches/fineweb_qa_pairs_{c}.jsonl", "w") as f: - for line in lines: - f.write(json.dumps(line) + "\n") - lines = [] - c += 1 - jsonl = get_json_request( - f"{c}_{i}", sample["text"], args.n_qa_pairs, args.gpt_model_name - ) - lines.append(jsonl) - + # ds = load_dataset( + # "parquet", + # data_files=f"./data/raw_datasets/fineweb_sharded/{args.shard_pattern}.parquet", + # split="train", + # streaming=True, + # ) + # os.makedirs("openai_batches", exist_ok=True) + # lines = [] + # c = 0 + # for i, sample in tqdm(enumerate(iter(ds))): + # if len(lines) >= 20_000: + # with open(f"openai_batches/fineweb_qa_pairs_{c}.jsonl", "w") as f: + # for line in lines: + # f.write(json.dumps(line) + "\n") + # lines = [] + # c += 1 + # jsonl = get_json_request( + # f"{c}_{i}", sample["text"], args.n_qa_pairs, args.gpt_model_name + # ) + # lines.append(jsonl) + # client = OpenAI() - - for file in glob("openai_batches/fineweb_qa_pairs_*.jsonl"): + res_files = glob("openai_batches/fineweb_qa_pairs_*_res.jsonl") + prompt_files = glob("openai_batches/fineweb_qa_pairs_*[!res].jsonl") + unprocessed_files = set(prompt_files) - set([f.replace("_res","") for f in res_files]) + for file in unprocessed_files: print(f"Submitting batch {file}") batch_input_file = client.files.create(file=open(file, "rb"), purpose="batch") batch_input_file_id = batch_input_file.id diff --git a/generate_one_big_fw_qa_req.py b/generate_one_big_fw_qa_req.py new file mode 100644 index 0000000..209bf63 --- /dev/null +++ b/generate_one_big_fw_qa_req.py @@ -0,0 +1,22 @@ +import json +from glob import glob + +if __name__ == "__main__": + res_files = glob("openai_batches/fineweb_qa_pairs_*_res.jsonl") + prompt_files = glob("openai_batches/fineweb_qa_pairs_*[!res].jsonl") + unprocessed_files = set(prompt_files) - set( + [f.replace("_res", "") for f in res_files] + ) + print(f"Found {len(unprocessed_files)} unprocessed files:\n{unprocessed_files}") + # Concatenate "body" from each line from unprocessed files + with open("openai_batches/fineweb_qa_pairs_large.jsonl", "w") as outfile: + for fname in unprocessed_files: + with open(fname) as infile: + for line in infile: + # Load the JSON object from the line + obj = json.loads(line) + # Extract the "body" and write it as a new line + if "body" in obj: + outfile.write(json.dumps(obj["body"]) + "\n") + + print(f"Concatenated 'body' from {len(unprocessed_files)} files") diff --git a/generate_qa_parallel.py b/generate_qa_parallel.py index 9455195..624cb67 100644 --- a/generate_qa_parallel.py +++ b/generate_qa_parallel.py @@ -467,8 +467,10 @@ if __name__ == "__main__": "--request_url", default="https://api.openai.com/v1/chat/completions" ) parser.add_argument("--api_key", default=os.getenv("OPENAI_API_KEY")) - parser.add_argument("--max_requests_per_minute", type=int, default=5_000 * 0.5) - parser.add_argument("--max_tokens_per_minute", type=int, default=2_000_000 * 0.8) + # tier 4 limit + parser.add_argument("--max_requests_per_minute", type=int, default=10_000 * 0.5) + parser.add_argument("--max_tokens_per_minute", type=int, default=10_000_000 * 0.5) + # parser.add_argument("--token_encoding_name", default="o200k_base") parser.add_argument("--max_attempts", type=int, default=5) parser.add_argument("--logging_level", default=logging.INFO) diff --git a/hotpot_qa_lengths.png b/hotpot_qa_lengths.png new file mode 100644 index 0000000..ea4cbd9 Binary files /dev/null and b/hotpot_qa_lengths.png differ diff --git a/hotpot_qa_lengths_filtered.png b/hotpot_qa_lengths_filtered.png new file mode 100644 index 0000000..63ddbf7 Binary files /dev/null and b/hotpot_qa_lengths_filtered.png differ diff --git a/openphi_prog_lengths.png b/openphi_prog_lengths.png new file mode 100644 index 0000000..9c2f399 Binary files /dev/null and b/openphi_prog_lengths.png differ diff --git a/openphi_prog_lengths_filtered.png b/openphi_prog_lengths_filtered.png new file mode 100644 index 0000000..c0e832a Binary files /dev/null and b/openphi_prog_lengths_filtered.png differ diff --git a/post_process_parallel_qa.py b/post_process_parallel_qa.py index 9ea7641..2d9483c 100644 --- a/post_process_parallel_qa.py +++ b/post_process_parallel_qa.py @@ -204,3 +204,9 @@ if __name__ == "__main__": print( f"Saved {fname}_neg.parquet to data/raw_datasets/ctx_qa/{fname}_neg.parquet" ) + # val_ds = ds.take(500) + # ds = ds.skip(500) + # ds.to_parquet(f"data/raw_datasets/fw_qa_large/{fname}.parquet") + # val_ds.to_parquet(f"data/raw_datasets/fw_qa_large/{fname}_val.parquet") + # print(f"Saved {fname} to data/raw_datasets/fw_qa_large/{fname}.parquet") + # print(f"Saved {fname}_val to data/raw_datasets/fw_qa_large/{fname}_val.parquet") diff --git a/run.sh b/run.sh new file mode 100755 index 0000000..6968e3b --- /dev/null +++ b/run.sh @@ -0,0 +1,9 @@ +job_id=$(sbatch --parsable batch.sh "$@") +while ! scontrol show job "$job_id" | grep -q "JobState=RUNNING"; do + echo "Waiting for job $job_id to start..." + sleep 5 +done + +echo "Job $job_id started" + +sattach $job_id.0 diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py new file mode 100644 index 0000000..e69de29 diff --git a/src/ctx_to_lora/lm_eval.py b/src/ctx_to_lora/lm_eval.py new file mode 100644 index 0000000..e69de29 diff --git a/src/ctx_to_lora/model_loading.py b/src/ctx_to_lora/model_loading.py index 915fdf5..173f07a 100644 --- a/src/ctx_to_lora/model_loading.py +++ b/src/ctx_to_lora/model_loading.py @@ -129,23 +129,29 @@ def get_model( is_vision_model = "Llama" in model_name_or_path and "Vision" in model_name_or_path if model_kwargs is not None: model_init_kwargs.update(model_kwargs) + + if is_vision_model: + # always use sdpa for vision models + model_init_kwargs["attn_implementation"] = "sdpa" + model_init_kwargs.pop("use_cache") + elif "gte" in model_name_or_path: + model_init_kwargs["torch_dtype"] = torch.float32 + model_init_kwargs.pop("use_cache") + if use_flash_attn: if ("gemma" not in model_name_or_path) and ("gte" not in model_name_or_path): model_init_kwargs["attn_implementation"] = "flash_attention_2" - if is_vision_model: - model_init_kwargs["attn_implementation"] = "sdpa" - model_init_kwargs.pop("use_cache") - # for training disable cache - elif "gte" in model_name_or_path: - model_init_kwargs["attn_implementation"] = "sdpa" - model_init_kwargs["torch_dtype"] = torch.float32 - model_init_kwargs.pop("use_cache") + elif "gte" in model_name_or_path: + model_init_kwargs["attn_implementation"] = "sdpa" + # if train and not is_vision_model: # model_init_kwargs["use_cache"] = False logger.debug(f"Model init kwargs: {model_init_kwargs}") if not is_vision_model: - # model = AutoModelForCausalLM.from_pretrained(**model_init_kwargs) - model = AutoModel.from_pretrained(**model_init_kwargs) + if "gte" in model_name_or_path: + model = AutoModel.from_pretrained(**model_init_kwargs) + else: + model = AutoModelForCausalLM.from_pretrained(**model_init_kwargs) else: model = MllamaForConditionalGeneration.from_pretrained(**model_init_kwargs) model = model.language_model