everything (before us migrate)

.gitignore

@@ -19,4 +19,5 @@ plots/
 *.pth
 *.bin
 *.safetensors
-*tfevents*
+*tfevents*
+watcher_state.yaml

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

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/

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
+

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

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

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

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")

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)

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")

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

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