apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-23 17:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

add per rank latent query + per layer hyper head

Browse source
This commit is contained in:
51616 2025-04-19 07:03:11 +00:00
parent 7866a43fa4
commit 49214d34df
39 changed files with 1277 additions and 391 deletions
Download patch file Download diff file Expand all files Collapse all files

29
batch.sh
View file

@ -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

30
batch_gemma_2_2b.sh
View file

@ -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/

31
batch_gemma_llama_new_up_and_down.sh
View file

@ -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

31
batch_gemma_llama_new_up_and_down_ln.sh
View file

@ -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

30
batch_gemma_llama_vision.sh
View file

@ -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

29
batch_llama_2_7b.sh
View file

@ -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

30
batch_llama_2b_.sh
View file

@ -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

29
batch_llama_3_8b_1k_latent.sh
View file

@ -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

61
configs/pretrain_all_2.yaml Normal file
View file

@ -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

39
dl_fineweb_edu.py Normal file
View file

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

232
generate_fw_edu_qa_vllm.py Normal file
View file

@ -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()

210
generate_fw_qa_vllm.py
View file

@ -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()

2
intx_sft.py
View file

@ -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()

36
process_fineweb.py
View file

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

3
requirements.txt
View file

@ -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

31
scripts/repeat_text/gemma_llama_xl.sh Normal file
View file

@ -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

31
scripts/scale_modules_fixed_latents/gemma_llama_xl.sh Normal file
View file

@ -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

31
scripts/scale_modules_fixed_latents/gemma_llama_xl_down_ln.sh Normal file
View file

@ -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

31
scripts/scale_modules_fixed_latents/gemma_llama_xl_ln_only.sh Normal file
View file

@ -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

31
scripts/scale_modules_fixed_latents/gemma_llama_xl_up_down.sh Normal file
View file

@ -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

20
batch_gemma_llama_new.sh → scripts/scale_modules_fixed_latents/gemma_llama_xl_up_down_ln.sh
View file

@ -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

31
scripts/scale_modules_no_lightweight/gemma_llama_2.sh Normal file
View file

@ -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

31
scripts/scale_modules_no_lightweight/gemma_llama_xl.sh Normal file
View file

@ -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

31
scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen.sh Normal file
View file

@ -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

32
scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen_latent_fac_1.sh Normal file
View file

@ -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

33
scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen_latent_fac_1_llama.sh Normal file
View file

@ -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

32
scripts/scale_modules_no_lightweight/gemma_llama_xl_per_rank_gen_latent_fac_1_r_16.sh Normal file
View file

@ -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

31
scripts/scale_modules_no_lightweight/gemma_llama_xl_r_16.sh Normal file
View file

@ -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

31
scripts/scale_modules_no_lightweight/gemma_llama_xl_r_16_ln.sh Normal file
View file

@ -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

31
scripts/scale_modules_no_lightweight/gemma_llama_xl_up_down.sh Normal file
View file

@ -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

31
scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl.sh Normal file
View file

@ -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

31
scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_down_ln.sh Normal file
View file

@ -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

31
scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_ln_only.sh Normal file
View file

@ -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

31
scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_up_down.sh Normal file
View file

@ -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

31
scripts/scale_modules_scale_latents_no_lightweight/gemma_llama_xl_up_down_ln.sh Normal file
View file

@ -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

4
src/ctx_to_lora/configs.py
View file

@ -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

45
src/ctx_to_lora/data_utils.py
View file

@ -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):

16
src/ctx_to_lora/eval.py
View file

@ -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,

168
src/ctx_to_lora/modeling_utils.py
View file

@ -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