mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
refactor after distillation (#6)
This commit is contained in:
Rujikorn Charakorn
GitHub
parent
2b824ed813
commit
610750e6fb
20 changed files with 694 additions and 1648 deletions
54
README.md
54
README.md
|
|
@ -31,12 +31,12 @@ print(model.decode(outputs))
|
|||
## 🏋️ Training
|
||||
### 🔢 HyperLoRA w/ context_numbers_10
|
||||
```bash
|
||||
WANDB_MODE=disabled run uv run intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=1 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --load_best_model_at_end=False --add_negative_prompt=False --add_repeat_prompt=False --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208
|
||||
WANDB_MODE=disabled run uv run intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=1 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208
|
||||
```
|
||||
|
||||
KL loss
|
||||
```bash
|
||||
WANDB_MODE=disabled run uv run intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=1 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --load_best_model_at_end=False --add_negative_prompt=False --add_repeat_prompt=False --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.0 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=True --eval_steps=99999999 --use_liger_kernel=False
|
||||
WANDB_MODE=disabled run uv run intx_sft.py configs/context_numbers_10_self_gen.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=2 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=100 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=2048 --max_packed_ctx_len=2048 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=True
|
||||
```
|
||||
### FWQA-v2 Level-0 Tiny
|
||||
```bash
|
||||
|
|
@ -76,6 +76,11 @@ uv run python data/generate_fav_num.py
|
|||
uv run python data/generate_fav_num_big.py
|
||||
```
|
||||
|
||||
Self-gen for the number toy dataset
|
||||
```bash
|
||||
run uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names context_numbers_2_10 --split train
|
||||
```
|
||||
|
||||
### SQuAD
|
||||
```bash
|
||||
# for some reason download directly through `load_dataset` does not work
|
||||
|
|
@ -145,7 +150,10 @@ uv run data/download_fineweb_edu.py
|
|||
1. Recursively generate more data! (depends on step 0)
|
||||
```bash
|
||||
# run from 000 to 0013
|
||||
run uv run python data/generate_fw_edu_qa_v2.py --shard_pattern "000_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048; run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it; run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_1*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it; run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_2*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it
|
||||
run uv run python data/generate_fw_edu_qa_v2.py --shard_pattern "000_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048;
|
||||
run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
|
||||
run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_1*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
|
||||
run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_2*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it
|
||||
```
|
||||
|
||||
2. Self-generated response QA data (depends on step 0 and 1)
|
||||
|
|
@ -157,47 +165,11 @@ uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_
|
|||
# val split
|
||||
uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_v2/min_0_to_2000/*_level_0_val.parquet'
|
||||
|
||||
# self-gen data for other ds listed in qa_no_fw.yaml
|
||||
uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/self_gen_qa_short_ctx_no_fw_qa.yaml
|
||||
# self-gen data for other ds listed in qa_short_ctx_self_gen_no_fw_qa.yaml
|
||||
uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_short_ctx_self_gen_no_fw_qa.yaml
|
||||
```
|
||||
|
||||
|
||||
<!-- ### Data Generation
|
||||
```bash
|
||||
# [WIP] there are other datasets where we used openAI to gen QA pairs
|
||||
|
||||
# 0. download fineweb_edu to `data/raw_datasets/fineweb_edu
|
||||
uv run data/download_fineweb_edu.py
|
||||
|
||||
# ===
|
||||
# 1 and 2 can be run in parallel (depends on step 0)
|
||||
|
||||
# 1. generate QA data (3 QAs for each context)
|
||||
# 000_0000[0-1] to 000_0000[8-9] for small split
|
||||
vllm_model=mistralai/Mistral-Small-3.1-24B-Instruct-2503 uv run python generate_fw_edu_qa_vllm.py "00*_*" 3
|
||||
|
||||
# 2. Augmenting pre-training data w/ paraphrasing
|
||||
# (000_00000 to 000_00009 for small split)
|
||||
vllm_model=mistralai/Mistral-Small-3.1-24B-Instruct-2503 uv run python data/generate_fw_edu_augment_vllm.py 000_00000
|
||||
|
||||
# ===
|
||||
# 3 and 4 can be run in parallel (depends on step 0 and 1)
|
||||
|
||||
# 3. Pre-training data (no GPU needed)
|
||||
# augment each context by concat'ing QAs generated in step 1 at the end
|
||||
uv run python data/generate_pretrain_from_fw_qa.py
|
||||
|
||||
# 4. self-generate SFT data
|
||||
# each context has 3 chat instances as training data
|
||||
# the responses are not the same as the ones in the pretraining though
|
||||
# Example commands using gemma-2-2b-it
|
||||
# self-gen data for fw_qa
|
||||
# 000_0000[0-1] to 000_0000[8-9] for small split
|
||||
uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_3/00*_*'
|
||||
# self-gen data for other ds listed in qa_no_fw.yaml
|
||||
uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_no_fw.yaml
|
||||
``` -->
|
||||
|
||||
<!--
|
||||
### GSM8k LoRA
|
||||
```bash
|
||||
|
|
|
|||
|
|
@ -2,8 +2,6 @@ output_dir: "" # just a placeholder
|
|||
bf16: true
|
||||
model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
|
||||
label_names: ["labels"]
|
||||
add_repeat_prompt: false
|
||||
add_negative_prompt: false
|
||||
# eval_on_start: True
|
||||
# eval_strategy: "steps"
|
||||
# eval_steps: 500
|
||||
|
|
|
|||
|
|
@ -2,8 +2,6 @@ output_dir: "" # just a placeholder
|
|||
bf16: true
|
||||
model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
|
||||
label_names: ["labels"]
|
||||
add_repeat_prompt: false
|
||||
add_negative_prompt: false
|
||||
# eval_on_start: True
|
||||
# eval_strategy: "steps"
|
||||
# eval_steps: 500
|
||||
|
|
|
|||
|
|
@ -39,11 +39,11 @@ target_modules:
|
|||
|
||||
# data
|
||||
train_ds_names:
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_3_tiny
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0/squad_compact
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0/pwc_compact
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0/ropes_compact
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0/drop_compact
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/ds_013*level_3*.parquet
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
|
||||
|
||||
val_ds_names:
|
||||
- squad
|
||||
|
|
|
|||
|
|
@ -239,23 +239,16 @@ def self_generate(
|
|||
|
||||
shard_name = "_" + os.path.basename(parquet_file).replace(".parquet", "")
|
||||
ds = load_dataset(path="parquet", data_files=[parquet_file], split="train")
|
||||
processing_fn = get_preprocessing_fn(ds_name, is_eval=False, is_pretrain=False)
|
||||
processing_fn = get_preprocessing_fn(ds_name, is_eval=False)
|
||||
ds = ds.map(processing_fn, num_proc=8)
|
||||
|
||||
else:
|
||||
ds_name = ds_name.split("/")[-1] # Extract just the dataset name
|
||||
|
||||
print(f"Loading dataset: {ds_name} with split: {split}")
|
||||
kwargs = dict(
|
||||
ds_name=ds_name,
|
||||
split=split,
|
||||
add_negative_prompt=False,
|
||||
add_repeat_prompt=False,
|
||||
repeat_prob=0,
|
||||
is_pretrain=False,
|
||||
)
|
||||
kwargs = dict(ds_name=ds_name, split=split)
|
||||
|
||||
ds = load_and_process_dataset(**kwargs, streaming=False, num_proc=8)
|
||||
ds = load_and_process_dataset(**kwargs, num_proc=8)
|
||||
print(f"Loaded dataset: {ds_name} with split: {split}")
|
||||
|
||||
if args.debug:
|
||||
|
|
|
|||
105
intx_sft.py
105
intx_sft.py
|
|
@ -32,7 +32,6 @@ from ctx_to_lora.configs import (
|
|||
)
|
||||
from ctx_to_lora.data.collator import ( # train_packed_collator,; DefaultDataCollator,
|
||||
flatten_if_not_packed,
|
||||
train_collator,
|
||||
)
|
||||
from ctx_to_lora.data.processing import get_tokenized_dataset, pack
|
||||
from ctx_to_lora.metrics import (
|
||||
|
|
@ -44,7 +43,6 @@ from ctx_to_lora.metrics import (
|
|||
)
|
||||
from ctx_to_lora.model_loading import (
|
||||
get_lora_config,
|
||||
get_model,
|
||||
get_model_and_tokenizer,
|
||||
get_tokenizer,
|
||||
)
|
||||
|
|
@ -125,6 +123,10 @@ def main():
|
|||
]
|
||||
)
|
||||
|
||||
assert ctx_args.use_sequence_packing, (
|
||||
f"Please set use_sequence_packing=True in {ctx_args}. It's faster!"
|
||||
)
|
||||
|
||||
set_seed(training_args.seed)
|
||||
checkpoint_dir = training_args.resume_from_checkpoint
|
||||
if checkpoint_dir and not os.path.isdir(checkpoint_dir):
|
||||
|
|
@ -218,8 +220,6 @@ def main():
|
|||
base_model,
|
||||
hypernet_config,
|
||||
ctx_encoder_args,
|
||||
ctx_args.use_kl_loss,
|
||||
ctx_args.use_sequence_packing,
|
||||
)
|
||||
|
||||
else:
|
||||
|
|
@ -230,7 +230,6 @@ def main():
|
|||
torch.load(ctx_args.from_pretrained_checkpoint, weights_only=False),
|
||||
train=True,
|
||||
use_flash_attn=model_args.use_flash_attn,
|
||||
use_sequence_packing=ctx_args.use_sequence_packing,
|
||||
)
|
||||
tokenizer = get_tokenizer(model.base_model.config.name_or_path, train=True)
|
||||
ctx_name = model.ctx_encoder_args.ctx_encoder_model_name_or_path
|
||||
|
|
@ -267,12 +266,6 @@ def main():
|
|||
logger.info("Loading dataset...")
|
||||
|
||||
add_ctx_to_chat = not isinstance(model, ModulatedPretrainedModel)
|
||||
tokenizer_kwargs = {"max_length": ctx_args.max_base_len} # not used
|
||||
ctx_tokenizer_kwargs = {"max_length": ctx_args.max_ctx_len} # not used for now
|
||||
|
||||
base_model = get_model(
|
||||
**vars(model_args), train=False, requires_grad=False, peft_config=None
|
||||
)
|
||||
|
||||
_get_tokenized_dataset = partial(
|
||||
get_tokenized_dataset,
|
||||
|
|
@ -280,18 +273,10 @@ def main():
|
|||
max_qas_per_sample=ctx_args.max_qas_per_sample,
|
||||
base_model_max_len=model.base_model.config.max_position_embeddings,
|
||||
tokenizer=tokenizer,
|
||||
tokenizer_kwargs=tokenizer_kwargs,
|
||||
ctx_model_max_len=model.ctx_encoder.config.max_position_embeddings,
|
||||
ctx_tokenizer=ctx_tokenizer,
|
||||
ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
|
||||
add_ctx_to_chat=add_ctx_to_chat,
|
||||
add_repeat_prompt=ctx_args.add_repeat_prompt,
|
||||
repeat_prob=ctx_args.repeat_prob,
|
||||
add_negative_prompt=ctx_args.add_negative_prompt,
|
||||
use_kl_loss=ctx_args.use_kl_loss,
|
||||
set_format=None if ctx_args.use_sequence_packing else "pt",
|
||||
base_model=base_model,
|
||||
# streaming=data_args.streaming,
|
||||
)
|
||||
tokenized_ds = {"train": dict(), "validation": dict()}
|
||||
for split, ds_names in zip(
|
||||
|
|
@ -300,7 +285,6 @@ def main():
|
|||
):
|
||||
if not ds_names:
|
||||
continue
|
||||
streaming = (split == "train") and data_args.streaming
|
||||
ctx_mgr = (
|
||||
training_args.main_process_first()
|
||||
if split == "train"
|
||||
|
|
@ -311,7 +295,7 @@ def main():
|
|||
# then other replicas can just load the cached dataset
|
||||
# we dont save cache for validation ds
|
||||
for ds_name in ds_names:
|
||||
ds = _get_tokenized_dataset(ds_name, split, streaming=streaming)
|
||||
ds = _get_tokenized_dataset(ds_name, split)
|
||||
|
||||
base_name = os.path.basename(ds_name)
|
||||
if ds_name.startswith("self_gen/"):
|
||||
|
|
@ -337,41 +321,8 @@ def main():
|
|||
val_indices = np.random.permutation(len(ds))[:n_val_samples]
|
||||
val_ds[ds_name] = val_ds[ds_name].select(val_indices)
|
||||
|
||||
# if data_args.streaming:
|
||||
# max_steps = ceil(
|
||||
# sum(DS_LEN[ds] for ds in train_ds)
|
||||
# * training_args.num_train_epochs
|
||||
# / training_args.per_device_train_batch_size
|
||||
# / training_args.gradient_accumulation_steps
|
||||
# / training_args.world_size
|
||||
# )
|
||||
# training_args.max_steps = max_steps
|
||||
|
||||
# # interleaving streaming datasets
|
||||
# # simplify the probs for smaller datasets
|
||||
# # slightly upsample those datasets
|
||||
# probs = [
|
||||
# 0.01 if "fw_qa" not in ds_name else 1 + 0.01 - 0.01 * len(train_ds)
|
||||
# for ds_name in train_ds
|
||||
# ]
|
||||
# train_ds = interleave_datasets(
|
||||
# list(train_ds.values()),
|
||||
# probabilities=probs,
|
||||
# stopping_strategy="all_exhausted",
|
||||
# seed=training_args.seed,
|
||||
# )
|
||||
# else:
|
||||
|
||||
train_ds_len = [len(ds) for ds in train_ds.values()]
|
||||
total_len = sum(train_ds_len)
|
||||
# max_steps = ceil(
|
||||
# total_len
|
||||
# * training_args.num_train_epochs
|
||||
# / training_args.per_device_train_batch_size
|
||||
# / training_args.gradient_accumulation_steps
|
||||
# / training_args.world_size
|
||||
# )
|
||||
# training_args.max_steps = max_steps
|
||||
train_ds = interleave_datasets(
|
||||
list(train_ds.values()),
|
||||
probabilities=get_ds_prob(train_ds_len, total_len),
|
||||
|
|
@ -380,35 +331,27 @@ def main():
|
|||
)
|
||||
logger.info(f"Train dataset length: {len(train_ds)}")
|
||||
|
||||
if ctx_args.use_sequence_packing:
|
||||
with training_args.main_process_first():
|
||||
logging.info("Packing dataset")
|
||||
old_ds_len = len(train_ds)
|
||||
train_ds = pack(
|
||||
train_ds,
|
||||
ctx_args.max_packed_inp_len,
|
||||
ctx_args.max_packed_ctx_len,
|
||||
max_packed_size=-1,
|
||||
num_proc=8,
|
||||
)
|
||||
logger.info(f"Train dataset length: {len(train_ds)}")
|
||||
logger.info(
|
||||
f"Avg. # of samples per packed sequence: {old_ds_len / len(train_ds)}"
|
||||
)
|
||||
logger.info("Setting per_device_train_batch_size to 1")
|
||||
training_args.per_device_train_batch_size = 1
|
||||
with training_args.main_process_first():
|
||||
logging.info("Packing dataset")
|
||||
old_ds_len = len(train_ds)
|
||||
train_ds = pack(
|
||||
train_ds,
|
||||
ctx_args.max_packed_inp_len,
|
||||
ctx_args.max_packed_ctx_len,
|
||||
max_packed_size=-1,
|
||||
num_proc=8,
|
||||
)
|
||||
logger.info(f"Train dataset length: {len(train_ds)}")
|
||||
logger.info(
|
||||
f"Avg. # of samples per packed sequence: {old_ds_len / len(train_ds)}"
|
||||
)
|
||||
logger.info("Setting per_device_train_batch_size to 1")
|
||||
training_args.per_device_train_batch_size = 1
|
||||
|
||||
logger.info(f"train_ds: {train_ds}")
|
||||
logger.info(f"val_ds: {val_ds}")
|
||||
|
||||
collator = (
|
||||
flatten_if_not_packed
|
||||
# DefaultDataCollator(return_tensors="pt")
|
||||
if ctx_args.use_sequence_packing
|
||||
else partial(train_collator, tokenizer=tokenizer)
|
||||
)
|
||||
|
||||
# TODO: use SFTTrainer instead? https://huggingface.co/docs/trl/en/sft_trainer
|
||||
collator = flatten_if_not_packed
|
||||
|
||||
# HACK [local patch]: deepspeed model loading problem (for resume training)
|
||||
# see https://github.com/microsoft/DeepSpeed/pull/6626/files
|
||||
|
|
@ -469,12 +412,10 @@ if __name__ == "__main__":
|
|||
os.environ["TOKENIZERS_PARALLELISM"] = "true"
|
||||
os.environ["WANDB_DIR"] = ".wandb/"
|
||||
os.environ["WANDB_PROJECT"] = "ctx_to_lora"
|
||||
os.environ["WANDB_WATCH"] = "" # "all"
|
||||
os.environ["WANDB_WATCH"] = ""
|
||||
os.environ["WANDB_CONSOLE"] = "off"
|
||||
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
||||
# os.environ["KMP_AFFINITY"] = "disabled" # fixing iterable dataset stuck
|
||||
os.environ["OMP_NUM_THREADS"] = "23"
|
||||
# os.environ["HF_DATASETS_IN_MEMORY_MAX_SIZE"] = "137438953472" # 128 TB
|
||||
torch._dynamo.config.capture_scalar_outputs = True
|
||||
|
||||
# DOES THIS CAUSE ACCURACY DEGRADATION?
|
||||
|
|
|
|||
20
run_eval.py
20
run_eval.py
|
|
@ -62,16 +62,16 @@ if __name__ == "__main__":
|
|||
|
||||
eval_batch_size_gen = cli_args.pop("eval_batch_size_gen")
|
||||
eval_batch_size = cli_args.pop("eval_batch_size")
|
||||
# run_eval(
|
||||
# **cli_args,
|
||||
# # cli_args.checkpoint_path,
|
||||
# # cli_args.model_name_or_path,
|
||||
# # cli_args.eval_batch_size,
|
||||
# # args,
|
||||
# # split=cli_args.split,
|
||||
# eval_batch_size=eval_batch_size,
|
||||
# generative=False,
|
||||
# )
|
||||
run_eval(
|
||||
**cli_args,
|
||||
# cli_args.checkpoint_path,
|
||||
# cli_args.model_name_or_path,
|
||||
# cli_args.eval_batch_size,
|
||||
# args,
|
||||
# split=cli_args.split,
|
||||
eval_batch_size=eval_batch_size,
|
||||
generative=False,
|
||||
)
|
||||
run_eval(
|
||||
**cli_args,
|
||||
# cli_args.checkpoint_path,
|
||||
|
|
|
|||
|
|
@ -6,14 +6,14 @@
|
|||
#SBATCH --output=slurm_logs/%x-%j.out
|
||||
#SBATCH --error=slurm_logs/%x-%j.out
|
||||
|
||||
port=$((10000 + ($SLURM_JOBID % 50000)))
|
||||
port=$((29560 + ($SLURM_JOBID % 10)))
|
||||
echo "Using port: $port"
|
||||
|
||||
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
|
||||
uv run accelerate launch --main_process_port $port \
|
||||
--num_processes=4 --gpu_ids all intx_sft.py $1 \
|
||||
--model_name_or_path=google/gemma-2-2b-it \
|
||||
--num_train_epochs=5 \
|
||||
--num_train_epochs=3 \
|
||||
--per_device_train_batch_size=-1 \
|
||||
--gradient_accumulation_steps=8 \
|
||||
--per_device_eval_batch_size=64 \
|
||||
|
|
@ -29,13 +29,11 @@ uv run accelerate launch --main_process_port $port \
|
|||
--learning_rate=4e-5 \
|
||||
--lora_dropout=0.0 \
|
||||
--neftune_noise_alpha=5 \
|
||||
--add_negative_prompt=False \
|
||||
--add_repeat_prompt=False \
|
||||
--use_sequence_packing=True \
|
||||
--max_qas_len=2048 \
|
||||
--max_qas_per_sample=1 \
|
||||
--max_packed_inp_len=8192 \
|
||||
--max_packed_ctx_len=16384 \
|
||||
--max_packed_ctx_len=8192 \
|
||||
--per_rank_gen=True \
|
||||
--per_layer_processing=True \
|
||||
--gen_lora_l1_reg_coef=0.1 \
|
||||
|
|
|
|||
|
|
@ -6,14 +6,14 @@
|
|||
#SBATCH --output=slurm_logs/%x-%j.out
|
||||
#SBATCH --error=slurm_logs/%x-%j.out
|
||||
|
||||
port=$((10000 + ($SLURM_JOBID % 50000)))
|
||||
port=$((29560 + ($SLURM_JOBID % 10)))
|
||||
echo "Using port: $port"
|
||||
|
||||
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
|
||||
uv run accelerate launch --main_process_port $port \
|
||||
--num_processes=4 --gpu_ids all intx_sft.py $1 \
|
||||
--model_name_or_path=google/gemma-2-2b-it \
|
||||
--num_train_epochs=5 \
|
||||
--num_train_epochs=3 \
|
||||
--per_device_train_batch_size=-1 \
|
||||
--gradient_accumulation_steps=16 \
|
||||
--per_device_eval_batch_size=64 \
|
||||
|
|
@ -29,8 +29,6 @@ uv run accelerate launch --main_process_port $port \
|
|||
--learning_rate=4e-5 \
|
||||
--lora_dropout=0.0 \
|
||||
--neftune_noise_alpha=5 \
|
||||
--add_negative_prompt=False \
|
||||
--add_repeat_prompt=False \
|
||||
--use_sequence_packing=True \
|
||||
--max_qas_len=2048 \
|
||||
--max_qas_per_sample=1 \
|
||||
|
|
|
|||
42
scripts/short_ctx/gemma_qa_short_ctx_old_conf_l2l_distill.sh
Normal file
42
scripts/short_ctx/gemma_qa_short_ctx_old_conf_l2l_distill.sh
Normal file
|
|
@ -0,0 +1,42 @@
|
|||
#!/bin/bash
|
||||
#SBATCH --job-name=ctxlora
|
||||
#SBATCH --nodes=1
|
||||
#SBATCH --partition=a3
|
||||
#SBATCH --gpus=4
|
||||
#SBATCH --output=slurm_logs/%x-%j.out
|
||||
#SBATCH --error=slurm_logs/%x-%j.out
|
||||
|
||||
port=$((29560 + ($SLURM_JOBID % 10)))
|
||||
echo "Using port: $port"
|
||||
|
||||
# --gradient_accumulation_steps=8 --gradient_clipping=1.0
|
||||
uv run accelerate launch --main_process_port $port \
|
||||
--num_processes=4 --gpu_ids all intx_sft.py $1 \
|
||||
--model_name_or_path=google/gemma-2-2b-it \
|
||||
--num_train_epochs=3 \
|
||||
--per_device_train_batch_size=-1 \
|
||||
--gradient_accumulation_steps=16 \
|
||||
--per_device_eval_batch_size=64 \
|
||||
--target_modules=down_proj \
|
||||
--num_blocks=9 \
|
||||
--num_self_attn_per_block=0 \
|
||||
--ctx_encoder_type=per_layer_activations \
|
||||
--n_latent_queries=8 \
|
||||
--num_pre_head_layers=1 \
|
||||
--lora_r=8 \
|
||||
--eval_strategy=no \
|
||||
--eval_steps=1000 \
|
||||
--save_steps=1000 \
|
||||
--learning_rate=4e-5 \
|
||||
--lora_dropout=0.0 \
|
||||
--neftune_noise_alpha=5 \
|
||||
--use_sequence_packing=True \
|
||||
--max_qas_len=2048 \
|
||||
--max_qas_per_sample=1 \
|
||||
--max_packed_inp_len=4096 \
|
||||
--max_packed_ctx_len=4096 \
|
||||
--per_rank_gen=True \
|
||||
--per_layer_processing=True \
|
||||
--gen_lora_l1_reg_coef=0.1 \
|
||||
--logging_steps=50 \
|
||||
--use_kl_loss=True
|
||||
|
|
@ -284,7 +284,7 @@ class LoRAArguments:
|
|||
metadata={"help": ("LoRA R value.")},
|
||||
)
|
||||
lora_dropout: float | None = field(
|
||||
default=0.02,
|
||||
default=0.0,
|
||||
metadata={"help": ("LoRA dropout.")},
|
||||
)
|
||||
target_modules: list[str] | None = field(
|
||||
|
|
@ -355,20 +355,6 @@ class CtxTrainingArguments:
|
|||
default=None,
|
||||
metadata={"help": "Wandb notes for the experiment."},
|
||||
)
|
||||
add_repeat_prompt: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Whether to add repeat prompt to the dataset."},
|
||||
)
|
||||
repeat_prob: float = field(
|
||||
default=0.0,
|
||||
metadata={
|
||||
"help": "Probability that a context will be used with repeat prompts."
|
||||
},
|
||||
)
|
||||
add_negative_prompt: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Whether to add negative prompt to the dataset."},
|
||||
)
|
||||
use_kl_loss: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Whether to use KL loss."},
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ from transformers.data import (
|
|||
default_data_collator,
|
||||
)
|
||||
|
||||
from ctx_to_lora.utils import check_is_iterable, concat_list
|
||||
from ctx_to_lora.utils import check_is_iterable
|
||||
|
||||
flattener = DataCollatorWithFlattening()
|
||||
|
||||
|
|
@ -14,6 +14,7 @@ def flatten_if_not_packed(inp_list):
|
|||
# no padding
|
||||
sample = inp_list[0]
|
||||
n = len(inp_list)
|
||||
# training data is packed
|
||||
if "position_ids" in sample:
|
||||
if n == 1:
|
||||
n_queries = sample.pop("n_queries")
|
||||
|
|
@ -27,16 +28,10 @@ def flatten_if_not_packed(inp_list):
|
|||
|
||||
# for eval data (not packed) during training
|
||||
need_flatten = check_is_iterable(sample["input_ids"][0])
|
||||
if need_flatten:
|
||||
n_queries = torch.tensor([len(x["input_ids"]) for x in inp_list])
|
||||
tokens = dict()
|
||||
for k in ["input_ids", "attention_mask", "labels"]:
|
||||
tokens[k] = concat_list([x[k] for x in inp_list])
|
||||
assert not need_flatten, f"Validation data should not be nested."
|
||||
|
||||
packed_inputs = flattener(tokens, return_tensors="pt")
|
||||
else:
|
||||
n_queries = torch.ones(len(inp_list), dtype=torch.int32)
|
||||
packed_inputs = flattener(inp_list, return_tensors="pt")
|
||||
n_queries = torch.ones(len(inp_list), dtype=torch.int32)
|
||||
packed_inputs = flattener(inp_list, return_tensors="pt")
|
||||
|
||||
packed_inputs["n_queries"] = n_queries
|
||||
if "ctx_ids" in sample:
|
||||
|
|
@ -53,167 +48,74 @@ def flatten_if_not_packed(inp_list):
|
|||
return packed_inputs
|
||||
|
||||
|
||||
def train_collator(inp_list, tokenizer):
|
||||
# input is a list of tokenized sequences
|
||||
padding_kwargs = dict(
|
||||
padding=True,
|
||||
padding_side="right",
|
||||
pad_to_multiple_of=8,
|
||||
return_tensors="pt",
|
||||
)
|
||||
sample = inp_list[0]
|
||||
ctx_ids = None
|
||||
if "ctx_ids" in sample:
|
||||
# pad to the longest ctx_len in the batch
|
||||
# which can have a different length from the input_ids, attn_mask, labels
|
||||
|
||||
ctx_ids = [example.pop("ctx_ids") for example in inp_list]
|
||||
ctx_ids = torch.nn.utils.rnn.pad_sequence(
|
||||
ctx_ids,
|
||||
batch_first=True,
|
||||
padding_value=0,
|
||||
)
|
||||
ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
|
||||
ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
|
||||
ctx_attn_mask,
|
||||
batch_first=True,
|
||||
padding_value=0,
|
||||
)
|
||||
|
||||
need_flatten = check_is_iterable(sample["input_ids"][0])
|
||||
|
||||
if need_flatten:
|
||||
# needed for one-lora-multi-q
|
||||
n_queries = torch.tensor([len(x["input_ids"]) for x in inp_list])
|
||||
tokens = dict()
|
||||
for k in ["input_ids", "attention_mask", "labels"]:
|
||||
tokens[k] = concat_list([x[k] for x in inp_list])
|
||||
labels = tokens.pop("labels")
|
||||
padded_seq = tokenizer.pad(tokens, **padding_kwargs)
|
||||
else:
|
||||
# needed for eval data during training
|
||||
labels = [x.pop("labels") for x in inp_list]
|
||||
padded_seq = tokenizer.pad(inp_list, **padding_kwargs)
|
||||
n_queries = torch.ones(len(inp_list), dtype=torch.int32)
|
||||
|
||||
# hacky explicit padding since the labels are not padded by default
|
||||
labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"]
|
||||
labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels)
|
||||
out = {**padded_seq, "labels": labels, "n_queries": n_queries}
|
||||
if ctx_ids is not None:
|
||||
out["ctx_ids"] = ctx_ids
|
||||
out["ctx_attn_mask"] = ctx_attn_mask
|
||||
if "logprobs_vals" in sample:
|
||||
# for training with logits
|
||||
logprobs_vals = [x.pop("logprobs_vals") for x in inp_list]
|
||||
# logprobs_vals = tokenizer.pad(
|
||||
# {"input_ids": logprobs_vals}, **padding_kwargs
|
||||
# )["input_ids"]
|
||||
out["logprobs_vals"] = logprobs_vals
|
||||
out["logprobs_indices"] = [x.pop("logprobs_indices") for x in inp_list]
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def eval_collator(inp_list, tokenizer):
|
||||
# only used for teacher-forcing eval
|
||||
# input is a list of tokenized sequences
|
||||
padding_kwargs = dict(
|
||||
padding=True,
|
||||
padding_side="right",
|
||||
pad_to_multiple_of=8,
|
||||
return_tensors="pt",
|
||||
)
|
||||
padding_kwargs = dict(padding=True, padding_side="right", return_tensors="pt")
|
||||
|
||||
ctx_ids = None
|
||||
if "ctx_ids" in inp_list[0]:
|
||||
has_ctx_ids = "ctx_ids" in inp_list[0]
|
||||
if has_ctx_ids:
|
||||
# pad to the longest ctx_len in the batch
|
||||
# which can have a different length from the input_ids, attn_mask, labels
|
||||
|
||||
ctx_ids = [example.pop("ctx_ids") for example in inp_list]
|
||||
ctx_attn_mask = [torch.ones_like(x) for x in ctx_ids]
|
||||
ctx_ids = torch.nn.utils.rnn.pad_sequence(
|
||||
ctx_ids,
|
||||
batch_first=True,
|
||||
padding_value=0,
|
||||
)
|
||||
ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
|
||||
ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
|
||||
ctx_attn_mask,
|
||||
batch_first=True,
|
||||
padding_value=0,
|
||||
)
|
||||
|
||||
chat_ids = None
|
||||
if "chat_ids" in inp_list[0]:
|
||||
chat_ids = [x.pop("chat_ids") for x in inp_list]
|
||||
chat_ids = torch.nn.utils.rnn.pad_sequence(
|
||||
chat_ids,
|
||||
batch_first=True,
|
||||
padding_value=0,
|
||||
)
|
||||
chat_attn_mask = [x.pop("chat_attn_mask") for x in inp_list]
|
||||
chat_attn_mask = torch.nn.utils.rnn.pad_sequence(
|
||||
chat_attn_mask,
|
||||
batch_first=True,
|
||||
padding_value=0,
|
||||
)
|
||||
chat_labels = [x.pop("chat_labels") for x in inp_list]
|
||||
chat_labels = torch.nn.utils.rnn.pad_sequence(
|
||||
chat_labels,
|
||||
batch_first=True,
|
||||
padding_value=-100,
|
||||
)
|
||||
chat_labels = torch.where(chat_attn_mask == 0, -100, chat_labels)
|
||||
for inp in inp_list:
|
||||
inp["attention_mask"] = torch.ones_like(inp["input_ids"])
|
||||
|
||||
labels = [x.pop("labels") for x in inp_list]
|
||||
# need to pass the whole inp bc we also track the lengths (with specal keys)
|
||||
padded_seq = tokenizer.pad(inp_list, **padding_kwargs)
|
||||
|
||||
# hacky explicit padding since the labels are not padded by default
|
||||
labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"]
|
||||
labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels)
|
||||
out = {**padded_seq, "labels": labels}
|
||||
if ctx_ids is not None:
|
||||
|
||||
if has_ctx_ids:
|
||||
out["ctx_ids"] = ctx_ids
|
||||
out["ctx_attn_mask"] = ctx_attn_mask
|
||||
if chat_ids is not None:
|
||||
out["chat_ids"] = chat_ids
|
||||
out["chat_attn_mask"] = chat_attn_mask
|
||||
out["chat_labels"] = chat_labels
|
||||
# if "input_ids_len" in inp_list[0]:
|
||||
# out["input_ids_len"] = [x.pop("input_ids_len") for x in inp_list]
|
||||
# if "ctx_ids_len" in inp_list[0]:
|
||||
# out["ctx_ids_len"] = [x.pop("ctx_ids_len") for x in inp_list]
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def generation_collator(inp_list, tokenizer):
|
||||
padding_kwargs = dict(padding=True, padding_side="left", return_tensors="pt")
|
||||
input_ids = [x.pop("input_ids") for x in inp_list]
|
||||
attn_mask = [x.pop("attention_mask") for x in inp_list]
|
||||
labels = [x.pop("labels") for x in inp_list]
|
||||
for i, label in enumerate(labels):
|
||||
# we don't include the labels in the output during generation
|
||||
# remove the response tokens
|
||||
idx = np.argmax(label != -100)
|
||||
idx = max(1, idx)
|
||||
input_ids[i] = input_ids[i][:idx]
|
||||
attn_mask[i] = attn_mask[i][:idx]
|
||||
attn_mask = [torch.ones_like(x) for x in input_ids]
|
||||
|
||||
out = tokenizer.pad(
|
||||
{"input_ids": input_ids, "attention_mask": attn_mask}, **padding_kwargs
|
||||
)
|
||||
|
||||
# we don't include the labels in the output
|
||||
# since we don't want to compute the loss on the labels
|
||||
# during generation
|
||||
if "ctx_ids" in inp_list[0]:
|
||||
# pad to the longest ctx_len in the batch
|
||||
# which can have a different length from the input_ids, attn_mask, labels
|
||||
ctx_ids = [example.pop("ctx_ids") for example in inp_list]
|
||||
ctx_attn_mask = [torch.ones_like(x) for x in ctx_ids]
|
||||
ctx_ids = torch.nn.utils.rnn.pad_sequence(
|
||||
ctx_ids,
|
||||
batch_first=True,
|
||||
padding_value=0,
|
||||
)
|
||||
ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
|
||||
ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
|
||||
ctx_attn_mask,
|
||||
batch_first=True,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
# based on
|
||||
# https://github.com/MeetKai/functionary/blob/aa3dbdd65f7e388f2386622606bdfeec95c2b863/functionary/train/packing/packed_dataset.py
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import pprint
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
|
@ -77,17 +80,13 @@ def pack_data_points_by_length(
|
|||
|
||||
def pack_data_points_FA(
|
||||
batch: dict[str, any],
|
||||
need_flatten: bool,
|
||||
) -> dict[str, np.ndarray]:
|
||||
if not batch:
|
||||
raise ValueError("Batch is empty")
|
||||
|
||||
# Pre-allocate lists with known sizes
|
||||
total_ctx_len = sum(len(x) for x in batch["ctx_ids"])
|
||||
if not need_flatten:
|
||||
total_inp_len = sum(len(x) for x in batch["input_ids"])
|
||||
else:
|
||||
total_inp_len = sum(len(y) for x in batch["input_ids"] for y in x)
|
||||
total_inp_len = sum(len(y) for x in batch["input_ids"] for y in x)
|
||||
|
||||
ctx_ids = np.empty(total_ctx_len, dtype=np.long)
|
||||
ctx_position_ids = np.empty(total_ctx_len, dtype=np.long)
|
||||
|
|
@ -95,9 +94,9 @@ def pack_data_points_FA(
|
|||
position_ids = np.empty(total_inp_len, dtype=np.long)
|
||||
labels = np.empty(total_inp_len, dtype=np.long)
|
||||
|
||||
has_logits = "logprobs_vals" in batch
|
||||
has_logprobs = "logprobs_vals" in batch
|
||||
|
||||
if has_logits:
|
||||
if has_logprobs:
|
||||
sequences = zip(
|
||||
batch["input_ids"],
|
||||
batch["labels"],
|
||||
|
|
@ -116,29 +115,26 @@ def pack_data_points_FA(
|
|||
for sample in sequences:
|
||||
input_ids_b, labels_b = sample[:2]
|
||||
inp_start = offset
|
||||
if need_flatten:
|
||||
# compute position_ids for each sub-list in input_ids_b
|
||||
local_start = inp_start
|
||||
for ids_b in input_ids_b:
|
||||
position_ids[local_start : local_start + len(ids_b)] = np.arange(
|
||||
len(ids_b), dtype=np.long
|
||||
)
|
||||
local_start += len(ids_b)
|
||||
|
||||
input_ids_b = concat_list(input_ids_b)
|
||||
labels_b = concat_list(labels_b)
|
||||
# compute position_ids for each sub-list in input_ids_b
|
||||
local_start = inp_start
|
||||
for ids_b in input_ids_b:
|
||||
position_ids[local_start : local_start + len(ids_b)] = np.arange(
|
||||
len(ids_b), dtype=np.long
|
||||
)
|
||||
local_start += len(ids_b)
|
||||
|
||||
input_ids_b = concat_list(input_ids_b)
|
||||
labels_b = concat_list(labels_b)
|
||||
|
||||
inp_len = len(input_ids_b)
|
||||
inp_end = offset + inp_len
|
||||
|
||||
if not need_flatten:
|
||||
position_ids[inp_start:inp_end] = np.arange(inp_len, dtype=np.long)
|
||||
|
||||
input_ids[inp_start:inp_end] = input_ids_b
|
||||
labels[inp_start:inp_end] = labels_b
|
||||
offset += inp_len
|
||||
|
||||
if has_logits:
|
||||
if has_logprobs:
|
||||
logprobs_vals_b, logprobs_indices_b = sample[2:]
|
||||
logprobs_vals_b = concat_list(logprobs_vals_b)
|
||||
logprobs_indices_b = concat_list(logprobs_indices_b)
|
||||
|
|
@ -163,7 +159,7 @@ def pack_data_points_FA(
|
|||
"position_ids": position_ids,
|
||||
"labels": labels,
|
||||
}
|
||||
if has_logits:
|
||||
if has_logprobs:
|
||||
out["logprobs_vals"] = logprobs_vals
|
||||
out["logprobs_indices"] = logprobs_indices
|
||||
return out
|
||||
|
|
@ -174,21 +170,17 @@ def pack_batch(
|
|||
max_packed_inp_len: int,
|
||||
max_packed_ctx_len: int,
|
||||
max_packed_size: int = -1,
|
||||
metadata_path: str = "",
|
||||
) -> dict[str, any]:
|
||||
need_flatten = check_is_iterable(batch["input_ids"][0][0])
|
||||
assert need_flatten
|
||||
assert need_flatten, (
|
||||
f"Packing requires the input_ids to be nested "
|
||||
f"(allowing multiple QAs per sample), but got {batch['input_ids'][0]}"
|
||||
)
|
||||
|
||||
if need_flatten:
|
||||
n_queries = [len(x) for x in batch["input_ids"]]
|
||||
inp_lens = [[len(y) for y in x] for x in batch["input_ids"]]
|
||||
# batch = {k: concat_list([x for x in batch[k]]) for k in keys}
|
||||
# for k in ["input_ids", "attention_mask", "labels"]:
|
||||
# batch[k] = concat_list([x for x in batch[k]])
|
||||
else:
|
||||
n_queries = [1] * len(batch["input_ids"])
|
||||
inp_lens = [[len(x)] for x in batch["input_ids"]]
|
||||
n_queries = [len(x) for x in batch["input_ids"]]
|
||||
inp_lens = [[len(y) for y in x] for x in batch["input_ids"]]
|
||||
inp_count = len(inp_lens)
|
||||
# total_original_tokens = sum(inp_lens)
|
||||
if "ctx_ids" not in batch:
|
||||
raise ValueError("Batch must contain 'ctx_ids' and 'labels' keys")
|
||||
# we do not pad so we can just take the length of the tokens
|
||||
|
|
@ -214,8 +206,8 @@ def pack_batch(
|
|||
"labels": [],
|
||||
"n_queries": [],
|
||||
}
|
||||
has_logits = "logprobs_vals" in batch
|
||||
if has_logits:
|
||||
has_logprobs = "logprobs_vals" in batch
|
||||
if has_logprobs:
|
||||
packed_batch["logprobs_vals"] = []
|
||||
packed_batch["logprobs_indices"] = []
|
||||
|
||||
|
|
@ -231,19 +223,19 @@ def pack_batch(
|
|||
"input_ids": batch["input_ids"][start_idx:end_idx],
|
||||
"labels": batch["labels"][start_idx:end_idx],
|
||||
}
|
||||
if has_logits:
|
||||
if has_logprobs:
|
||||
group_items["logprobs_vals"] = batch["logprobs_vals"][start_idx:end_idx]
|
||||
group_items["logprobs_indices"] = batch["logprobs_indices"][
|
||||
start_idx:end_idx
|
||||
]
|
||||
packed_item = pack_data_points_FA(group_items, need_flatten)
|
||||
packed_item = pack_data_points_FA(group_items)
|
||||
packed_batch["ctx_ids"].append(packed_item["ctx_ids"])
|
||||
packed_batch["ctx_position_ids"].append(packed_item["ctx_position_ids"])
|
||||
packed_batch["input_ids"].append(packed_item["input_ids"])
|
||||
packed_batch["position_ids"].append(packed_item["position_ids"])
|
||||
packed_batch["labels"].append(packed_item["labels"])
|
||||
packed_batch["n_queries"].append(n_queries[start_idx:end_idx])
|
||||
if has_logits:
|
||||
if has_logprobs:
|
||||
packed_batch["logprobs_vals"].append(packed_item["logprobs_vals"])
|
||||
packed_batch["logprobs_indices"].append(packed_item["logprobs_indices"])
|
||||
|
||||
|
|
@ -272,18 +264,33 @@ def pack_batch(
|
|||
np.mean(ctx_packing_efficiency_ratios) if ctx_packing_efficiency_ratios else 0
|
||||
)
|
||||
|
||||
logging.debug(
|
||||
f"Packing stats - Original samples: {inp_count}\n"
|
||||
f"# Packed samples: {len(idx_pairs)}\n"
|
||||
f"Avg inp packing efficiency: {avg_inp_packing_efficiency:.3f}\n"
|
||||
f"Avg ctx packing efficiency: {avg_ctx_packing_efficiency:.3f}\n\n"
|
||||
f"Input IDs length stats:\n"
|
||||
f" Avg: {np.mean(packed_inp_lens_arr):.1f}, Std: {np.std(packed_inp_lens_arr):.1f}, "
|
||||
f"Min: {np.min(packed_inp_lens_arr)}, Max: {np.max(packed_inp_lens_arr)}\n"
|
||||
f"Context IDs length stats:\n"
|
||||
f" Avg: {np.mean(packed_ctx_lens_arr):.1f}, Std: {np.std(packed_ctx_lens_arr):.1f}, "
|
||||
f"Min: {np.min(packed_ctx_lens_arr)}, Max: {np.max(packed_ctx_lens_arr)}\n\n"
|
||||
)
|
||||
# Create packing statistics dictionary
|
||||
packing_stats = {
|
||||
"original_samples": inp_count,
|
||||
"packed_samples": len(idx_pairs),
|
||||
"avg_inp_packing_efficiency": float(avg_inp_packing_efficiency),
|
||||
"avg_ctx_packing_efficiency": float(avg_ctx_packing_efficiency),
|
||||
"input_ids_length_stats": {
|
||||
"avg": float(np.mean(packed_inp_lens_arr)),
|
||||
"std": float(np.std(packed_inp_lens_arr)),
|
||||
"min": int(np.min(packed_inp_lens_arr)),
|
||||
"max": int(np.max(packed_inp_lens_arr)),
|
||||
},
|
||||
"context_ids_length_stats": {
|
||||
"avg": float(np.mean(packed_ctx_lens_arr)),
|
||||
"std": float(np.std(packed_ctx_lens_arr)),
|
||||
"min": int(np.min(packed_ctx_lens_arr)),
|
||||
"max": int(np.max(packed_ctx_lens_arr)),
|
||||
},
|
||||
}
|
||||
|
||||
# Save to metadata_path if provided
|
||||
if metadata_path:
|
||||
os.makedirs(os.path.dirname(metadata_path), exist_ok=True)
|
||||
with open(metadata_path, "w") as f:
|
||||
json.dump(packing_stats, f, indent=4)
|
||||
|
||||
logging.debug(f"Packing stats:\n{pprint.pformat(packing_stats, indent=2)}")
|
||||
|
||||
return packed_batch
|
||||
|
||||
|
|
|
|||
280
src/ctx_to_lora/data/preprocessing_fn.py
Normal file
280
src/ctx_to_lora/data/preprocessing_fn.py
Normal file
|
|
@ -0,0 +1,280 @@
|
|||
import logging
|
||||
import random
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from ctx_to_lora.data.definitions import CLOSED_QA_INTX_TEMPLATES, EVAL_INTX_TEMPLATES
|
||||
from ctx_to_lora.utils import concat_list
|
||||
|
||||
logger = logging.getLogger()
|
||||
|
||||
|
||||
def closed_qa_prompting(prompt: str):
|
||||
template = random.choice(CLOSED_QA_INTX_TEMPLATES)
|
||||
return template.format(input=prompt)
|
||||
|
||||
|
||||
def get_preprocessing_fn(
|
||||
ds_name: str,
|
||||
is_eval: bool,
|
||||
) -> Callable[[dict[str, Any]], dict[str, Any]]:
|
||||
"""
|
||||
Get preprocessing function for a specific dataset.
|
||||
|
||||
Args:
|
||||
ds_name: Name of the dataset
|
||||
|
||||
Returns:
|
||||
A preprocessing function that takes and returns a dictionary
|
||||
"""
|
||||
f = lambda x: x
|
||||
if ds_name.startswith("self_gen") or ds_name.endswith("_compact"):
|
||||
# already processed data, do nothing
|
||||
return f
|
||||
|
||||
if "fw_qa_v2" in ds_name:
|
||||
|
||||
def f(sample):
|
||||
# get questions/answers from all levels in the ds
|
||||
q_cols = [col for col in sample.keys() if col.startswith("prompts_level")]
|
||||
r_cols = [col for col in sample.keys() if col.startswith("responses_level")]
|
||||
questions = concat_list([sample[col] for col in q_cols])
|
||||
responses = concat_list([sample[col] for col in r_cols])
|
||||
min_len = min(len(questions), len(responses))
|
||||
|
||||
if min_len == 0:
|
||||
return {
|
||||
"context": None,
|
||||
"prompts": None,
|
||||
"responses": None,
|
||||
}
|
||||
|
||||
return {
|
||||
"context": sample["context"],
|
||||
"prompts": questions[:min_len],
|
||||
"responses": responses[:min_len],
|
||||
}
|
||||
|
||||
elif ds_name.startswith("longbench"):
|
||||
|
||||
def f(sample):
|
||||
return {
|
||||
"context": sample["context"],
|
||||
"prompt": sample["input"],
|
||||
"response": sample["answers"][0],
|
||||
}
|
||||
|
||||
elif ds_name == "negative_nq":
|
||||
|
||||
def f(sample):
|
||||
q = sample["prompt"]
|
||||
prompt = closed_qa_prompting(q) if not is_eval else q
|
||||
return {
|
||||
"context": sample["context"],
|
||||
"prompt": prompt,
|
||||
"response": sample["answer"],
|
||||
}
|
||||
|
||||
elif ds_name == "triviaqa_retrieved":
|
||||
# only used for eval
|
||||
def f(sample):
|
||||
return {
|
||||
"context": sample["context"],
|
||||
"prompt": sample["prompt"],
|
||||
"response": sample["answer"],
|
||||
}
|
||||
|
||||
elif ds_name == "pwc" or ds_name == "pwc_tiny":
|
||||
# original pwc
|
||||
def f(sample):
|
||||
return {
|
||||
"context": sample["input"],
|
||||
"prompt": sample["prompt"],
|
||||
"response": sample["answer"],
|
||||
}
|
||||
|
||||
elif ds_name.startswith("hotpot_qa"):
|
||||
|
||||
def f(sample):
|
||||
txt = ""
|
||||
for p in sample["context"]["sentences"]:
|
||||
txt += " " + "".join(p)
|
||||
|
||||
q = sample["question"]
|
||||
prompt = closed_qa_prompting(q) if not is_eval else q
|
||||
|
||||
return {
|
||||
"context": txt.strip(),
|
||||
"prompt": prompt,
|
||||
"response": sample["answer"],
|
||||
}
|
||||
|
||||
elif ds_name == "squad":
|
||||
# original squad
|
||||
def f(sample):
|
||||
q = sample["question"]
|
||||
prompt = closed_qa_prompting(q) if not is_eval else q
|
||||
return {
|
||||
"context": sample["context"],
|
||||
"prompt": prompt,
|
||||
"response": sample["answers"]["text"][0],
|
||||
}
|
||||
|
||||
elif ds_name == "drop":
|
||||
|
||||
def f(sample):
|
||||
q = sample["question"]
|
||||
prompt = closed_qa_prompting(q) if not is_eval else q
|
||||
return {
|
||||
"context": sample["passage"],
|
||||
"prompt": prompt,
|
||||
"response": sample["answers_spans"]["spans"][0],
|
||||
}
|
||||
|
||||
elif ds_name == "ropes":
|
||||
ctx_template = "{background}\n{situation}"
|
||||
|
||||
def f(sample):
|
||||
response = sample["answers"]["text"][0]
|
||||
bg_txt = sample["background"]
|
||||
situation_txt = sample["situation"]
|
||||
ctx = ctx_template.format(background=bg_txt, situation=situation_txt)
|
||||
q = sample["question"]
|
||||
q = closed_qa_prompting(q) if not is_eval else q
|
||||
return {"context": ctx, "prompt": q, "response": response}
|
||||
|
||||
elif ds_name in ["narrativeqa", "quoref"]: # , "ropes"]:
|
||||
|
||||
def f(sample):
|
||||
response = sample["answers"][0]
|
||||
if isinstance(response, list):
|
||||
response = response[0]
|
||||
q = sample["messages"][-1]["content"]
|
||||
prompt = closed_qa_prompting(q) if not is_eval else q
|
||||
return {
|
||||
"context": sample["document"],
|
||||
"prompt": prompt,
|
||||
"response": response,
|
||||
}
|
||||
|
||||
elif ds_name == "synthetic_convqa":
|
||||
|
||||
def f(sample):
|
||||
response = sample["answers"][0]
|
||||
if isinstance(response, list):
|
||||
response = response[0]
|
||||
return {
|
||||
"context": sample["document"],
|
||||
"prompt": sample["messages"][-1]["content"],
|
||||
"response": response,
|
||||
}
|
||||
|
||||
elif ds_name == "booksum":
|
||||
prompt_templates = [
|
||||
"Summarization the provided text.",
|
||||
"# Summary",
|
||||
"### Summary",
|
||||
"Summary of the text",
|
||||
]
|
||||
# TODO: use these templates for training
|
||||
# analysis_templates = []
|
||||
# summary_templates = []
|
||||
|
||||
def f(sample):
|
||||
return {
|
||||
"context": sample["chapter"],
|
||||
"prompt": "Summarize the provided text.",
|
||||
"response": sample["summary_text"],
|
||||
}
|
||||
|
||||
elif ds_name == "gov_report":
|
||||
|
||||
def f(sample):
|
||||
return {
|
||||
"context": sample["report"],
|
||||
"prompt": "Summarize the provided text.",
|
||||
"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):
|
||||
return {
|
||||
"context": sample["problem"],
|
||||
"prompt": sample["problem"],
|
||||
"response": sample["generated_solution"],
|
||||
}
|
||||
|
||||
elif "gsm8k" in ds_name:
|
||||
|
||||
def f(sample):
|
||||
return {
|
||||
"context": sample["question"],
|
||||
"prompt": sample["question"],
|
||||
"response": sample["answer"],
|
||||
}
|
||||
|
||||
elif "opencoder-edu" in ds_name:
|
||||
|
||||
def f(sample):
|
||||
return {
|
||||
"context": sample["instruction"],
|
||||
"prompt": sample["instruction"],
|
||||
"response": "```python\n" + sample["code"].strip() + "\n```",
|
||||
}
|
||||
|
||||
if is_eval and (ds_name in EVAL_INTX_TEMPLATES):
|
||||
prompt_template = EVAL_INTX_TEMPLATES[ds_name]
|
||||
|
||||
def eval_intx_decorator(f):
|
||||
def g(sample):
|
||||
sample = f(sample)
|
||||
assert "prompt" in sample, (
|
||||
f"Expected 'prompt' in sample, got {sample.keys()}"
|
||||
)
|
||||
sample["prompt"] = prompt_template.format(input=sample["prompt"])
|
||||
return sample
|
||||
|
||||
return g
|
||||
|
||||
f = eval_intx_decorator(f)
|
||||
|
||||
def maybe_convert_to_list(f):
|
||||
def g(sample):
|
||||
sample = f(sample)
|
||||
if "prompt" in sample:
|
||||
sample["prompts"] = [sample.pop("prompt")]
|
||||
if "response" in sample:
|
||||
sample["responses"] = [sample.pop("response")]
|
||||
return sample
|
||||
|
||||
return g
|
||||
|
||||
f = maybe_convert_to_list(f)
|
||||
|
||||
if "self_gen" not in ds_name:
|
||||
|
||||
def strip_response(f):
|
||||
def g(sample):
|
||||
sample = f(sample)
|
||||
if "responses" in sample and bool(sample["responses"]):
|
||||
sample["responses"] = [
|
||||
r.strip() if isinstance(r, str) else r
|
||||
for r in sample["responses"]
|
||||
]
|
||||
return sample
|
||||
|
||||
return g
|
||||
|
||||
f = strip_response(f)
|
||||
|
||||
return f
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -680,6 +680,7 @@ def evaluate(
|
|||
train=False,
|
||||
base_model_kwargs=model_kwargs,
|
||||
use_flash_attn=True,
|
||||
use_sequence_packing=False, # for generation
|
||||
)
|
||||
else:
|
||||
model = get_model(
|
||||
|
|
@ -724,8 +725,6 @@ def evaluate(
|
|||
if ctx_tokenizer.pad_token_id is None:
|
||||
ctx_tokenizer.pad_token_id = ctx_tokenizer.eos_token_id
|
||||
|
||||
tokenizer_kwargs = {}
|
||||
ctx_tokenizer_kwargs = {}
|
||||
add_ctx_to_chat = (
|
||||
not isinstance(model, ModulatedPretrainedModel) and not args.remove_context
|
||||
)
|
||||
|
|
@ -741,13 +740,9 @@ def evaluate(
|
|||
max_qas_per_sample=1,
|
||||
base_model_max_len=model.base_model.config.max_position_embeddings,
|
||||
tokenizer=tokenizer,
|
||||
tokenizer_kwargs=tokenizer_kwargs,
|
||||
ctx_model_max_len=ctx_model_max_len,
|
||||
ctx_tokenizer=ctx_tokenizer,
|
||||
ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
|
||||
add_ctx_to_chat=add_ctx_to_chat,
|
||||
add_repeat_prompt=False,
|
||||
add_negative_prompt=False,
|
||||
use_kl_loss=False,
|
||||
set_format="pt",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -591,19 +591,18 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
base_model: PeftModel,
|
||||
hypernet_config: HypernetConfig,
|
||||
ctx_encoder_args: CtxEncoderArguments,
|
||||
# use_kl_loss is only used for training
|
||||
use_kl_loss: bool = False,
|
||||
use_sequence_packing: bool = False,
|
||||
use_base_input_as_ctx: bool = False,
|
||||
# need non-packed inputs for generation
|
||||
use_sequence_packing: bool = True,
|
||||
):
|
||||
assert not use_base_input_as_ctx
|
||||
super().__init__()
|
||||
self.device = base_model.device
|
||||
self.peft_config = base_model.peft_config["default"]
|
||||
self.hypernet_config = hypernet_config
|
||||
self.ctx_encoder_args = ctx_encoder_args
|
||||
self.use_kl_loss = use_kl_loss
|
||||
self.use_sequence_packing = use_sequence_packing
|
||||
self.use_base_input_as_ctx = use_base_input_as_ctx
|
||||
self.use_sequence_packing = use_sequence_packing
|
||||
self.lora_id = 1
|
||||
self.active_adapters = []
|
||||
|
||||
|
|
@ -618,6 +617,7 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
train: bool = True,
|
||||
base_model_kwargs: dict = None,
|
||||
use_flash_attn: bool = True,
|
||||
**kwargs: Any,
|
||||
):
|
||||
lora_config = state_dict["hypernet_config"].lora_config
|
||||
print(f"lora_config: {lora_config}")
|
||||
|
|
@ -634,7 +634,7 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
if getattr(hypernet_config, "num_pre_head_layers", None) is None:
|
||||
hypernet_config.num_pre_head_layers = 4
|
||||
ctx_encoder_args = state_dict["ctx_encoder_args"]
|
||||
model = cls(base_model, hypernet_config, ctx_encoder_args)
|
||||
model = cls(base_model, hypernet_config, ctx_encoder_args, **kwargs)
|
||||
model.load_state_dict(state_dict)
|
||||
return model
|
||||
|
||||
|
|
@ -650,7 +650,7 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
|
||||
layers = get_layers(self.base_model)
|
||||
lora_forward_fn = (
|
||||
lora_forward if not self.use_sequence_packing else lora_forward_packed
|
||||
lora_forward_packed if self.use_sequence_packing else lora_forward
|
||||
)
|
||||
for layer_idx in self.hypernet.layer_indices:
|
||||
for module_info in get_peft_modules(layers[layer_idx], self.peft_config):
|
||||
|
|
@ -909,44 +909,57 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
else:
|
||||
return model_outputs
|
||||
|
||||
# @torch.inference_mode()
|
||||
# def generate_with_multi_loras(
|
||||
# self,
|
||||
# ctx_ids: Integer[Tensor, "bs ctx_length"],
|
||||
# ctx_attn_mask: Integer[Tensor, "bs ctx_length"] | None = None,
|
||||
# ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None,
|
||||
# *model_inputs_args: Any,
|
||||
# **model_inputs_kwargs: dict[str, Any],
|
||||
# ):
|
||||
# # TODO: make this persistent
|
||||
# generated_loras, _ = self.generate_weights(
|
||||
# ctx_ids, ctx_attn_mask, ctx_position_ids
|
||||
# )
|
||||
# # sum loras
|
||||
# for lora_at_module in generated_loras.values():
|
||||
# if len(lora_at_module["A"]) > 1:
|
||||
# print("Multiple LoRAs are generated, summing them up")
|
||||
# lora_at_module["A"] = torch.mean(lora_at_module["A"], dim=0, keepdim=True)
|
||||
# lora_at_module["B"] = torch.mean(lora_at_module["B"], dim=0, keepdim=True)
|
||||
@torch.inference_mode()
|
||||
def generate_with_multi_loras(
|
||||
self,
|
||||
ctx_ids: Integer[Tensor, "bs ctx_length"],
|
||||
ctx_attn_mask: Integer[Tensor, "bs ctx_length"] | None = None,
|
||||
ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None,
|
||||
n_queries: Integer[Tensor, "n_ctx"] | None = None,
|
||||
*model_inputs_args: Any,
|
||||
**model_inputs_kwargs: dict[str, Any],
|
||||
):
|
||||
# TODO: make this persistent
|
||||
generated_loras, _ = self.generate_weights(
|
||||
ctx_ids, ctx_attn_mask, ctx_position_ids
|
||||
)
|
||||
# sum loras
|
||||
for lora_at_module in generated_loras.values():
|
||||
if len(lora_at_module["A"]) > 1:
|
||||
print("Multiple LoRAs are generated, summing them up")
|
||||
lora_at_module["A"] = torch.mean(lora_at_module["A"], dim=0, keepdim=True)
|
||||
lora_at_module["B"] = torch.mean(lora_at_module["B"], dim=0, keepdim=True)
|
||||
|
||||
# position_ids = (
|
||||
# model_inputs_kwargs["position_ids"]
|
||||
# if "position_ids" in model_inputs_kwargs
|
||||
# else None
|
||||
# )
|
||||
# apply lora hook to the base model
|
||||
position_ids = (
|
||||
model_inputs_kwargs["position_ids"]
|
||||
if "position_ids" in model_inputs_kwargs
|
||||
else None
|
||||
)
|
||||
if n_queries is None:
|
||||
if ctx_position_ids is None:
|
||||
n_queries = torch.ones(
|
||||
ctx_ids.shape[0], dtype=torch.int32, device=self.device
|
||||
)
|
||||
else:
|
||||
# quite redundant (we do cu_seqlens many places)
|
||||
# TODO: compute cu_seqlens here and propagate that
|
||||
n_queries = torch.ones(
|
||||
(ctx_position_ids == 0).sum(), dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
# # apply loras
|
||||
# with apply_generated_loras(
|
||||
# self.base_model,
|
||||
# generated_loras,
|
||||
# self.hypernet.layer_indices,
|
||||
# position_ids,
|
||||
# self.training,
|
||||
# ):
|
||||
# model_outputs = self.base_model.generate(
|
||||
# *model_inputs_args, **model_inputs_kwargs
|
||||
# )
|
||||
# return model_outputs
|
||||
apply_lora_to_layers(
|
||||
self.base_model,
|
||||
self.hypernet.layer_indices,
|
||||
generated_loras,
|
||||
n_queries,
|
||||
position_ids, # TODO: remove (cant be used with generation?)
|
||||
)
|
||||
|
||||
model_outputs = self.base_model.generate(
|
||||
*model_inputs_args, **model_inputs_kwargs
|
||||
)
|
||||
return model_outputs
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(
|
||||
|
|
@ -1033,93 +1046,6 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
return model_outputs
|
||||
|
||||
|
||||
# class ModulatedModelWithSharedInput(nn.Module):
|
||||
# def __init__(
|
||||
# self,
|
||||
# modulated_model: ModulatedPretrainedModel,
|
||||
# base_tokenizer: PreTrainedTokenizerBase,
|
||||
# ctx_tokenizer: PreTrainedTokenizerBase,
|
||||
# ctx_end_predicate: str | None = None,
|
||||
# remove_ctx_from_base_input: bool = False,
|
||||
# ):
|
||||
# super().__init__()
|
||||
# self.modulated_model = modulated_model
|
||||
# self.base_tokenizer = base_tokenizer
|
||||
# self.ctx_tokenizer = ctx_tokenizer
|
||||
# self.register_module("modulated_model", self.modulated_model)
|
||||
# self.ctx_end_predicate = ctx_end_predicate
|
||||
# self.remove_ctx_from_base_input = remove_ctx_from_base_input
|
||||
|
||||
# # delegate to self.modulated_model
|
||||
# @property
|
||||
# def config(self):
|
||||
# return self.modulated_model.config
|
||||
|
||||
# @property
|
||||
# def device(self):
|
||||
# return self.modulated_model.device
|
||||
|
||||
# def tie_weights(self):
|
||||
# self.modulated_model.base_model.tie_weights()
|
||||
|
||||
# def state_dict(self, *args, **kwargs):
|
||||
# return self.modulated_model.state_dict(*args, **kwargs)
|
||||
|
||||
# def forward(self, *args, **kwargs):
|
||||
# input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0]
|
||||
# input_txts = self.base_tokenizer.batch_decode(
|
||||
# input_ids,
|
||||
# skip_special_tokens=True,
|
||||
# clean_up_tokenization_spaces=True,
|
||||
# )
|
||||
# ctx_txts = deepcopy(input_txts)
|
||||
# if self.ctx_end_predicate:
|
||||
# ctx_txts = [
|
||||
# txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts
|
||||
# ]
|
||||
# ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to(
|
||||
# self.device
|
||||
# )
|
||||
# ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask
|
||||
# if self.remove_ctx_from_base_input:
|
||||
# raise NotImplementedError("Not implemented")
|
||||
# # input_txts = [
|
||||
# # txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts
|
||||
# # ]
|
||||
# # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True)
|
||||
# # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask
|
||||
# # kwargs["input_ids"] = input_ids
|
||||
# # kwargs["attention_mask"] = attention_mask
|
||||
# return self.modulated_model(ctx_ids, ctx_attn_mask, *args, **kwargs)
|
||||
|
||||
# def generate(self, *args, **kwargs):
|
||||
# input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0]
|
||||
# input_txts = self.base_tokenizer.batch_decode(
|
||||
# input_ids,
|
||||
# skip_special_tokens=True,
|
||||
# clean_up_tokenization_spaces=True,
|
||||
# )
|
||||
# ctx_txts = deepcopy(input_txts)
|
||||
# if self.ctx_end_predicate:
|
||||
# ctx_txts = [
|
||||
# txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts
|
||||
# ]
|
||||
# ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to(
|
||||
# self.device
|
||||
# )
|
||||
# ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask
|
||||
# if self.remove_ctx_from_base_input:
|
||||
# raise NotImplementedError("Not implemented")
|
||||
# # input_txts = [
|
||||
# # txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts
|
||||
# # ]
|
||||
# # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True)
|
||||
# # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask
|
||||
# # kwargs["input_ids"] = input_ids
|
||||
# # kwargs["attention_mask"] = attention_mask
|
||||
# return self.modulated_model.generate(ctx_ids, ctx_attn_mask, *args, **kwargs)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def apply_generated_layernorm(
|
||||
base_model: nn.Module,
|
||||
|
|
|
|||
|
|
@ -59,7 +59,6 @@ def lora_forward_packed(
|
|||
# repeated_A = A.repeat_interleave(seq_lens, dim=0, output_size=seq_lens.sum())
|
||||
# # [tot_seq_len, d_out, r]
|
||||
# repeated_B = B.repeat_interleave(seq_lens, dim=0, output_size=seq_lens.sum())
|
||||
# TODO: is ther a better way to implement this?
|
||||
repeated_A = A.repeat_interleave(n_qs, dim=0, output_size=tot_q)
|
||||
repeated_A = repeated_A.repeat_interleave(seq_lens, dim=0, output_size=tot_len)
|
||||
|
||||
|
|
|
|||
|
|
@ -13,25 +13,6 @@ from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
|
|||
logger = logging.getLogger()
|
||||
|
||||
|
||||
def add_pre_label_token(labels, is_packed):
|
||||
"""
|
||||
Add a token before the first label token.
|
||||
1 is used as a placeholder for the first label token.
|
||||
We compute the loss against the teacher logits anyway.
|
||||
"""
|
||||
if is_packed:
|
||||
non_label_mask = torch.where(labels == -100, True, False)
|
||||
diff = torch.diff(labels, append=torch.tensor([[-100]], device=labels.device))
|
||||
pre_label_pos = torch.where(non_label_mask * (diff != 0))
|
||||
labels[pre_label_pos] = 1
|
||||
else:
|
||||
bsz = labels.shape[0]
|
||||
first_label_pos = torch.argmax((labels != -100).float(), axis=-1)
|
||||
labels[torch.arange(bsz), first_label_pos - 1] = 1
|
||||
return labels
|
||||
|
||||
|
||||
# TODO: test
|
||||
class DistillationTrainer(Trainer):
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
|
||||
|
|
@ -46,71 +27,61 @@ class DistillationTrainer(Trainer):
|
|||
if num_items_in_batch is not None:
|
||||
loss_kwargs["num_items_in_batch"] = num_items_in_batch
|
||||
inputs = {**inputs, **loss_kwargs}
|
||||
|
||||
labels = inputs.pop("labels", None) # remove labels if present
|
||||
# NOTE: have to account that we save +1 # token before the first label token
|
||||
is_packed = "position_ids" in inputs
|
||||
# labels = add_pre_label_token(labels, is_packed)
|
||||
label_pos = torch.where(labels != -100)
|
||||
use_kl = "logprobs_vals" in inputs
|
||||
# TODO: handle when val data doesnt have logprobs
|
||||
if use_kl:
|
||||
labels = inputs.pop("labels", None)
|
||||
label_pos = torch.where(labels != -100)
|
||||
# TODO: add efficient logits indexing?
|
||||
# https://github.com/huggingface/trl/pull/2773
|
||||
# + use `num_logits_to_keep` to limit the number of logits
|
||||
outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True)
|
||||
target_logp = inputs.pop("logprobs_vals")
|
||||
indices = inputs.pop("logprobs_indices")
|
||||
|
||||
# packed
|
||||
if isinstance(target_logp, torch.Tensor):
|
||||
target_logp = target_logp.squeeze(0)
|
||||
indices = indices.squeeze(0)
|
||||
# non-packed max_qas_per_sample=1
|
||||
elif isinstance(target_logp[0], torch.Tensor):
|
||||
target_logp = torch.cat([x[0] for x in target_logp], dim=0)
|
||||
indices = torch.cat([x[0] for x in indices], dim=0)
|
||||
# non-packed max_qas_per_sample>1
|
||||
if not use_kl:
|
||||
loss = outputs["loss"]
|
||||
|
||||
else:
|
||||
# convert to tensor
|
||||
raise NotImplementedError(
|
||||
"Non-packed data with max_qas_per_sample > 1 is not supported yet."
|
||||
target_logp = inputs.pop("logprobs_vals")
|
||||
indices = inputs.pop("logprobs_indices")
|
||||
|
||||
# packed
|
||||
if isinstance(target_logp, torch.Tensor):
|
||||
target_logp = target_logp.squeeze(0)
|
||||
indices = indices.squeeze(0)
|
||||
# non-packed max_qas_per_sample=1
|
||||
elif isinstance(target_logp[0], torch.Tensor):
|
||||
target_logp = torch.cat([x[0] for x in target_logp], dim=0)
|
||||
indices = torch.cat([x[0] for x in indices], dim=0)
|
||||
# non-packed max_qas_per_sample>1
|
||||
else:
|
||||
# convert to tensor
|
||||
raise NotImplementedError(
|
||||
"Non-packed data with max_qas_per_sample > 1 is not supported yet."
|
||||
)
|
||||
|
||||
assert label_pos[0].shape[0] == target_logp.shape[0], (
|
||||
"Label positions and target log probabilities should have the same # tokens."
|
||||
f"Got : {label_pos[0].shape[0]=} and {target_logp.shape[0]=}"
|
||||
)
|
||||
|
||||
assert label_pos[0].shape[0] == target_logp.shape[0], (
|
||||
"Label positions and target log probabilities should have the same # tokens."
|
||||
f"Got : {label_pos[0].shape[0]=} and {target_logp.shape[0]=}"
|
||||
)
|
||||
# TODO: per-ctx average loss?
|
||||
##### KL loss
|
||||
outputs_logits = outputs.logits[
|
||||
label_pos[0], label_pos[1] - 1
|
||||
] # shift back 1
|
||||
|
||||
# TODO: per-ctx average loss?
|
||||
##### KL loss
|
||||
outputs_logits = outputs.logits[label_pos[0], label_pos[1] - 1] # shift back 1
|
||||
teacher_logp = torch.full_like(outputs_logits, -torch.inf)
|
||||
teacher_logp.scatter_(1, indices, target_logp)
|
||||
reduction = "batchmean" if num_items_in_batch is None else "sum"
|
||||
|
||||
teacher_logp = torch.full_like(outputs_logits, -torch.inf)
|
||||
teacher_logp.scatter_(1, indices, target_logp)
|
||||
reduction = "batchmean" if num_items_in_batch is None else "sum"
|
||||
|
||||
# loss = torch.nn.functional.kl_div(
|
||||
# torch.nn.functional.log_softmax(
|
||||
# outputs_logits / self.distillation_temperature, dim=-1
|
||||
# ),
|
||||
# torch.nn.functional.softmax(
|
||||
# teacher_logits / self.distillation_temperature, dim=-1
|
||||
# ),
|
||||
# reduction=reduction,
|
||||
# log_target=False,
|
||||
# ) * (self.distillation_temperature**2)
|
||||
# loss = loss / num_items_in_batch if num_items_in_batch is not None else loss
|
||||
|
||||
# p = torch.nn.functional.softmax(
|
||||
# teacher_logits / self.distillation_temperature, dim=-1
|
||||
# )
|
||||
p = teacher_logp.exp()
|
||||
logq = nn.functional.log_softmax(outputs_logits, dim=-1)
|
||||
|
||||
loss = -torch.sum(p * logq, dim=-1) * (self.distillation_temperature**2)
|
||||
if reduction == "batchmean":
|
||||
loss = loss.mean()
|
||||
elif reduction == "sum" and num_items_in_batch is not None:
|
||||
loss = loss.sum() / num_items_in_batch
|
||||
#####
|
||||
p = teacher_logp.exp()
|
||||
logq = nn.functional.log_softmax(outputs_logits, dim=-1)
|
||||
loss = -torch.sum(p * logq, dim=-1) * (self.distillation_temperature**2)
|
||||
if reduction == "batchmean":
|
||||
loss = loss.mean()
|
||||
elif reduction == "sum" and num_items_in_batch is not None:
|
||||
loss = loss.sum() / num_items_in_batch
|
||||
#####
|
||||
|
||||
##### unpack gen lora dict and compute regularization loss
|
||||
l1_norm = 0
|
||||
|
|
@ -121,7 +92,8 @@ class DistillationTrainer(Trainer):
|
|||
self.args.logging_strategy == IntervalStrategy.STEPS
|
||||
and self.state.global_step % self.state.logging_steps == 0
|
||||
):
|
||||
self.log({"ce_loss": loss.item(), "gen_lora_l1_norm": l1_norm.item()})
|
||||
loss_field = "kl_loss" if use_kl else "ce_loss"
|
||||
self.log({loss_field: loss.item(), "gen_lora_l1_norm": l1_norm.item()})
|
||||
loss += self.gen_lora_l1_reg_coef * l1_norm
|
||||
#####
|
||||
|
||||
|
|
|
|||
|
|
@ -502,6 +502,7 @@ def load_checkpoint():
|
|||
state_dict,
|
||||
train=False,
|
||||
use_flash_attn=True,
|
||||
use_sequence_packing=False, # for generation
|
||||
)
|
||||
modulated_model = modulated_model.to(device).to(torch.bfloat16)
|
||||
modulated_model.eval()
|
||||
|
|
@ -547,7 +548,7 @@ def process_multiple_contexts(contexts, ctx_tokenizer):
|
|||
ctx_ids = [
|
||||
torch.tensor(ctx_id, dtype=torch.long, device=device) for ctx_id in ctx_ids
|
||||
]
|
||||
ctx_attn_mask = tokenized_contexts["ctx_attn_mask"]
|
||||
ctx_attn_mask = [torch.ones_like(ids) for ids in ctx_ids]
|
||||
ctx_attn_mask = [
|
||||
torch.tensor(ctx_attn_mask, dtype=torch.long, device=device)
|
||||
for ctx_attn_mask in ctx_attn_mask
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue