From 178e7165ea280a649987461129b8fa47b020c1a5 Mon Sep 17 00:00:00 2001 From: 51616 Date: Sat, 9 Aug 2025 15:20:37 +0000 Subject: [PATCH] per-rank bias + skip interleave if cached+ add bnb --- README.md | 32 +-- ...ctx_self_gen_lv1_closed_1_and_lv3_l2l.yaml | 42 ++++ ...hort_ctx_self_gen_lv1_closed_qa_1_l2l.yaml | 14 +- pyproject.toml | 1 + .../gemma_qa_short_ctx_exp_default_8_gpus.sh | 29 +++ src/ctx_to_lora/configs.py | 8 + src/ctx_to_lora/data/processing.py | 35 ++- src/ctx_to_lora/model_loading.py | 18 ++ src/ctx_to_lora/modeling/aggregator.py | 2 +- src/ctx_to_lora/modeling/hypernet.py | 218 ++++++++++-------- train.py | 47 +--- uv.lock | 16 ++ 12 files changed, 302 insertions(+), 160 deletions(-) create mode 100644 configs/medium_exp/qa_short_ctx_self_gen_lv1_closed_1_and_lv3_l2l.yaml create mode 100644 scripts/short_ctx/gemma_qa_short_ctx_exp_default_8_gpus.sh diff --git a/README.md b/README.md index c8f59a4..1e6bf5a 100644 --- a/README.md +++ b/README.md @@ -47,28 +47,8 @@ WANDB_MODE=disabled run uv run train.py configs/fw_qa_v2_level_0_tiny.yaml --mod ```bash WANDB_MODE=disabled uv run train.py configs/squad.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=64 --gradient_accumulation_steps=8 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=1 --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 --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --max_packed_inp_len=16000 --max_packed_ctx_len=32000 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=10 ``` - - - ### Favourite numbers data generation ```bash uv run python data/generate_fav_num.py @@ -88,6 +68,18 @@ run uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --d HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download --repo-type dataset rajpurkar/squad --local-dir data/raw_datasets/squad ``` +### Self-gen data upload/download +```bash +# create a bucket (needed only once) +# gsutil mb -l EU gs://ctx-to-lora/ + +# uploading from login node to gcp bucket +gsutil -m rsync -r data/raw_datasets/self_gen gs://ctx-to-lora/data/raw_datasets/self_gen + +# downloading from gcp bucket to login node +gsutil -m rsync -r gs://ctx-to-lora/data/raw_datasets/self_gen data/raw_datasets/self_gen +``` + ### Data Generation (v2) ***Download a subset of generated Fineweb-QA directly*** ```bash diff --git a/configs/medium_exp/qa_short_ctx_self_gen_lv1_closed_1_and_lv3_l2l.yaml b/configs/medium_exp/qa_short_ctx_self_gen_lv1_closed_1_and_lv3_l2l.yaml new file mode 100644 index 0000000..4ce1d5d --- /dev/null +++ b/configs/medium_exp/qa_short_ctx_self_gen_lv1_closed_1_and_lv3_l2l.yaml @@ -0,0 +1,42 @@ +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +use_kl_loss: true + +ctx_encoder_type: per_layer_activations +n_latent_queries: 8 +num_blocks: 9 +num_self_attn_per_block: 0 + +gradient_accumulation_steps: 11 +max_packed_inp_len: 6144 +max_packed_ctx_len: 6144 + + + + +# data +train_ds_names: + - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet + - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*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 + - pwc + - drop + - ropes + # - hotpot_qa + # - negative_nq + # - booksum + # - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet + - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet + # - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet + # - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet + diff --git a/configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_l2l.yaml b/configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_l2l.yaml index 151fc95..cf18365 100644 --- a/configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_l2l.yaml +++ b/configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_l2l.yaml @@ -11,9 +11,9 @@ n_latent_queries: 8 num_blocks: 9 num_self_attn_per_block: 0 -gradient_accumulation_steps: 16 -max_packed_inp_len: 4096 -max_packed_ctx_len: 4096 +gradient_accumulation_steps: 11 +max_packed_inp_len: 6144 +max_packed_ctx_len: 6144 @@ -31,11 +31,11 @@ val_ds_names: - pwc - drop - ropes - - hotpot_qa - - negative_nq - - booksum + # - hotpot_qa + # - negative_nq + # - booksum # - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet # - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet - - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet + # - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet diff --git a/pyproject.toml b/pyproject.toml index 5ca7a74..2ca24e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dependencies = [ "opt-einsum>=3.4.0", "kagglehub[hf-datasets]>=0.3.12", "kaggle>=1.7.4.5", + "bitsandbytes>=0.46.1", ] [build-system] diff --git a/scripts/short_ctx/gemma_qa_short_ctx_exp_default_8_gpus.sh b/scripts/short_ctx/gemma_qa_short_ctx_exp_default_8_gpus.sh new file mode 100644 index 0000000..5685bd3 --- /dev/null +++ b/scripts/short_ctx/gemma_qa_short_ctx_exp_default_8_gpus.sh @@ -0,0 +1,29 @@ +#!/bin/bash +#SBATCH --job-name=ctxlora +#SBATCH --nodes=1 +#SBATCH --partition=a3 +#SBATCH --gpus=8 +#SBATCH --output=slurm_logs/%x-%j.out +#SBATCH --error=slurm_logs/%x-%j.out + +port=$((10000 + ($SLURM_JOBID % 50000))) +echo "Using port: $port" + +# Default arguments - these can be overridden by passing arguments to this script +default_args=( + "--model_name_or_path=google/gemma-2-2b-it" + "--max_steps=50_000" + "--target_modules=down_proj" + "--lora_r=8" + "--eval_strategy=no" + "--max_qas_len=2048" + "--max_qas_per_sample=1" + "--per_rank_gen=True" + "--per_layer_processing=True" + "--gen_lora_l1_reg_coef=0.1" +) + +# Pass all script arguments to the training command +# $1 comes first, then defaults, then remaining arguments (which can override defaults) +uv run accelerate launch --main_process_port $port \ +--num_processes=8 --gpu_ids all train.py $1 "${default_args[@]}" "${@:2}" diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index 11be386..b1832c8 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -207,6 +207,7 @@ class TrainingArguments(TrainingArguments): default=0.01, metadata={"help": "Weight decay for the optimizer."}, ) + # TODO: is paged_adam_w faster? (bigger batch size -> fewer grad accum steps) optim: str = field( default="adamw_torch_fused", metadata={"help": "Optimizer."}, @@ -293,6 +294,10 @@ class TrainingArguments(TrainingArguments): default=True, metadata={"help": "Whether to log the first step."}, ) + ddp_find_unused_parameters: bool = field( + default=False, + metadata={"help": "Whether to find unused parameters in DDP."}, + ) ddp_timeout: int = field( default=2**20, metadata={"help": "Timeout for distributed data parallel training."}, @@ -470,6 +475,9 @@ class HypernetArguments: default=False, metadata={"help": "Whether to use per-rank generation."}, ) + use_per_rank_bias: bool = field( + default=False, metadata={"help": "Whether to use per-rank bias."} + ) per_layer_processing: bool = field( default=False, metadata={"help": "Whether to use per-layer processing (after preceiver)."}, diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index c787889..2e17bad 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -5,12 +5,13 @@ import random from collections.abc import Callable from glob import glob from hashlib import sha256 +from math import isclose from os import path from typing import Any import datasets import torch -from datasets import Dataset, is_caching_enabled, load_dataset +from datasets import Dataset, interleave_datasets, is_caching_enabled, load_dataset from transformers import PreTrainedTokenizerBase from ctx_to_lora.data.definitions import ( @@ -50,6 +51,24 @@ COLS_TO_KEEP_TOKENIZED = [ ] +def get_ds_prob(train_ds_len: list[int], total_len: int): + # if a dataset is smaller than 1%, make it 1% + probs = [0 for _ in train_ds_len] + for i, ds_len in enumerate(train_ds_len): + if ds_len / total_len <= 0.01: + probs[i] = 0.01 + res_probs = 1 - sum(probs) + res_total_len = sum([l for l in train_ds_len if (l / total_len) > 0.01]) + for i, ds_len in enumerate(train_ds_len): + if (ds_len / total_len) > 0.01: + probs[i] = ds_len / res_total_len * res_probs + logger.debug(f"Dataset probabilities: {probs}") + assert isclose(sum(probs), 1.0), ( + f"Probs sum to {sum(probs)} ({probs}), expected 1.0" + ) + return probs + + def load_answers(ds_name, split): if ds_name.startswith("longbench"): @@ -759,10 +778,11 @@ def tokenize_ctx_text( def pack( - ds: Dataset, + ds_dict: dict[str, Dataset], max_packed_inp_len: int, max_packed_ctx_len: int, max_packed_size: int, + seed: int, num_proc: int = 0, ): # this would generate another cache file for the already concat'd + packed ds @@ -771,7 +791,9 @@ def pack( max_packed_ctx_len=max_packed_ctx_len, max_packed_size=max_packed_size, ) - ds_hash = sha256((ds._fingerprint + json.dumps(kwargs)).encode()).hexdigest() + sorted_keys = sorted(ds_dict) + ds_fingerprint = "|".join([ds_dict[k]._fingerprint for k in sorted_keys]) + ds_hash = sha256((ds_fingerprint + json.dumps(kwargs)).encode()).hexdigest() ds_path = f"{TRANSFORMED_DATA_DIR}/packed_{ds_hash}" logger.info( f"Packing ds {ds_hash} with {max_packed_inp_len=} and {max_packed_ctx_len=}" @@ -780,6 +802,13 @@ def pack( logger.info(f"Loading a cached packed dataset for {ds_path}") return datasets.load_from_disk(ds_path) else: + train_ds = interleave_datasets( + list(train_ds.values()), + probabilities=get_ds_prob(train_ds_len, total_len), + seed=seed, + stopping_strategy="all_exhausted", + ) + logger.info(f"Train dataset length: {len(train_ds)}") ds = ds.map( pack_batch, fn_kwargs={ diff --git a/src/ctx_to_lora/model_loading.py b/src/ctx_to_lora/model_loading.py index 456562a..672ad56 100644 --- a/src/ctx_to_lora/model_loading.py +++ b/src/ctx_to_lora/model_loading.py @@ -9,6 +9,7 @@ from transformers import ( AutoModel, AutoModelForCausalLM, AutoTokenizer, + BitsAndBytesConfig, Gemma3ForConditionalGeneration, ) @@ -33,6 +34,7 @@ def get_model_and_tokenizer( peft_config=None, model_kwargs=None, tokenizer_kwargs=None, + use_q_lora=False, device="cuda", dtype=torch.bfloat16, ): @@ -43,6 +45,7 @@ def get_model_and_tokenizer( use_flash_attn, peft_config, model_kwargs, + use_q_lora, device, dtype, ) @@ -95,6 +98,7 @@ def get_model( use_flash_attn=True, peft_config=None, model_kwargs=None, + use_q_lora=False, device="cuda", dtype=torch.bfloat16, ): @@ -128,6 +132,20 @@ def get_model( model_init_kwargs["torch_dtype"] = torch.float32 model_init_kwargs.pop("use_cache") + if use_q_lora: + # https://huggingface.co/blog/4bit-transformers-bitsandbytes + # https://colab.research.google.com/drive/1VoYNfYDKcKRQRor98Zbf2-9VQTtGJ24k?usp=sharing + # see bitsandbytes for the quantization implementation https://github.com/bitsandbytes-foundation/bitsandbytes + # see unsloth https://huggingface.co/docs/trl/v0.7.11/en/sft_trainer#accelerate-fine-tuning-2x-using-unsloth + # does work currently bc it modifies the forward pass call of Linear + bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + ) + model_init_kwargs["quantization_config"] = bnb_config + logger.debug(f"Model init kwargs: {model_init_kwargs}") if not is_vision_model: if is_bidir_model: diff --git a/src/ctx_to_lora/modeling/aggregator.py b/src/ctx_to_lora/modeling/aggregator.py index 7594bee..3a472bb 100644 --- a/src/ctx_to_lora/modeling/aggregator.py +++ b/src/ctx_to_lora/modeling/aggregator.py @@ -199,7 +199,7 @@ class Perceiver(nn.Module): per_layer_size = self.num_modules * self.r + self.num_extra_modules x = rearrange( x, - ("(num_layers bs) (per_layer_sz) d -> bs (num_layers per_layer_sz) d"), + "(num_layers bs) (per_layer_sz) d -> bs (num_layers per_layer_sz) d", num_layers=self.num_layers, per_layer_sz=per_layer_size, ) diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index 1e32560..d2daff8 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/src/ctx_to_lora/modeling/hypernet.py @@ -66,6 +66,7 @@ class HypernetConfig: use_light_weight_lora: bool light_weight_latent_size: int per_rank_gen: bool + use_per_rank_bias: bool per_layer_processing: bool use_token_mixing: bool num_pre_head_layers: int @@ -353,98 +354,135 @@ class HyperLoRA(nn.Module): if self.config.per_rank_gen: if n_modules == 1: if self.config.per_layer_processing: - self.head = Mix( - "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", - weight_shape="n_layers d_latent d_lora", - # bias_shape=None, # no bias - bias_shape="n_layers 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="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, - ) + if self.config.use_per_rank_bias: + self.head = Mix( + "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", + weight_shape="n_layers d_latent d_lora", + # bias_shape=None, # no bias + bias_shape="n_layers r 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_layers d_latent d_lora", + # bias_shape=None, # no bias + bias_shape="n_layers 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: + # if self.config.use_per_rank_bias: + # 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="r 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="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: if self.config.per_layer_processing: - self.head = Mix( - "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", - weight_shape="n_layers n_modules d_latent d_lora", - # bias_shape=None, # no bias - bias_shape="n_layers 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: - 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: - if n_modules == 1: - if self.config.per_layer_processing: - self.head = Mix( - "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", - weight_shape="n_layers d_latent r d_lora", - # bias_shape=None, # no bias - bias_shape="n_layers r d_lora", - 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 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 - if self.config.per_layer_processing: - self.head = Mix( - "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", - weight_shape="n_layers n_modules d_latent r d_lora", - # bias_shape=None, # no bias - bias_shape="n_layers 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, - ) - else: - 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.config.use_per_rank_bias: + self.head = Mix( + "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", + weight_shape="n_layers n_modules d_latent d_lora", + # bias_shape=None, # no bias + bias_shape="n_layers n_modules r 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: + self.head = Mix( + "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", + weight_shape="n_layers n_modules d_latent d_lora", + # bias_shape=None, # no bias + bias_shape="n_layers 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: + # 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: + # if n_modules == 1: + # if self.config.per_layer_processing: + # self.head = Mix( + # "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", + # weight_shape="n_layers d_latent r d_lora", + # # bias_shape=None, # no bias + # bias_shape="n_layers r d_lora", + # 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 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 + # if self.config.per_layer_processing: + # self.head = Mix( + # "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", + # weight_shape="n_layers n_modules d_latent r d_lora", + # # bias_shape=None, # no bias + # bias_shape="n_layers 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, + # ) + # else: + # 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( diff --git a/train.py b/train.py index cdf8d82..2f5c0b7 100755 --- a/train.py +++ b/train.py @@ -3,15 +3,11 @@ import logging import os from copy import deepcopy from functools import partial -from math import isclose import numpy as np import torch import wandb -from datasets import ( - disable_caching, - interleave_datasets, -) +from datasets import disable_caching from peft import PeftModel from transformers import ( AutoConfig, @@ -66,24 +62,6 @@ logger = logging.getLogger() LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) -def get_ds_prob(train_ds_len: list[int], total_len: int): - # if a dataset is smaller than 1%, make it 1% - probs = [0 for _ in train_ds_len] - for i, ds_len in enumerate(train_ds_len): - if ds_len / total_len <= 0.01: - probs[i] = 0.01 - res_probs = 1 - sum(probs) - res_total_len = sum([l for l in train_ds_len if (l / total_len) > 0.01]) - for i, ds_len in enumerate(train_ds_len): - if (ds_len / total_len) > 0.01: - probs[i] = ds_len / res_total_len * res_probs - logger.debug(f"Dataset probabilities: {probs}") - assert isclose(sum(probs), 1.0), ( - f"Probs sum to {sum(probs)} ({probs}), expected 1.0" - ) - return probs - - def main(): ############ Argument parsing parser = ArgumentParser( @@ -183,6 +161,7 @@ def main(): train=True, requires_grad=False, # ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE, peft_config=get_lora_config(model_name, **vars(lora_args)), + # use_q_lora=True, ) ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path if ctx_name is not None: @@ -328,31 +307,21 @@ def main(): val_indices = np.random.permutation(len(ds))[:n_val_samples] val_ds[ds_name] = val_ds[ds_name].select(val_indices) - train_ds_len = [len(ds) for ds in train_ds.values()] - total_len = sum(train_ds_len) - - # TODO: cache this? - train_ds = interleave_datasets( - list(train_ds.values()), - probabilities=get_ds_prob(train_ds_len, total_len), - seed=training_args.seed, - stopping_strategy="all_exhausted", - ) - logger.info(f"Train dataset length: {len(train_ds)}") - + total_samples = sum([len(ds) for ds in train_ds.values()]) + logging.info(f"Total samples before packing: {total_samples}") + logging.info("Packing dataset") 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, + seed=training_args.seed, num_proc=16, ) - logger.info(f"Train dataset length: {len(train_ds)}") + logger.info(f"Packed dataset length: {len(train_ds)}") logger.info( - f"Avg. # of samples per packed sequence: {old_ds_len / len(train_ds)}" + f"Avg. # of samples per packed sequence: {total_samples / len(train_ds)}" ) logger.info("Setting per_device_train_batch_size to 1") training_args.per_device_train_batch_size = 1 diff --git a/uv.lock b/uv.lock index 9a5afbe..2da0f61 100644 --- a/uv.lock +++ b/uv.lock @@ -308,6 +308,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/50/cd/30110dc0ffcf3b131156077b90e9f60ed75711223f306da4db08eff8403b/beautifulsoup4-4.13.4-py3-none-any.whl", hash = "sha256:9bbbb14bfde9d79f38b8cd5f8c7c85f4b8f2523190ebed90e950a8dea4cb1c4b", size = 187285, upload-time = "2025-04-15T17:05:12.221Z" }, ] +[[package]] +name = "bitsandbytes" +version = "0.46.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "torch" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/d2/b2/9dadb4f8dca3948e35c1ebfee75ca82353e41468b41ff785430595f8e6f0/bitsandbytes-0.46.1-py3-none-manylinux_2_24_aarch64.whl", hash = "sha256:21b349f776d04c6c1380405961081de29c84f49640b79d3d199b6d719818da84", size = 30713241, upload-time = "2025-07-02T19:44:21.857Z" }, + { url = "https://files.pythonhosted.org/packages/6b/1e/c26dbcb46cebb49fa6b17ff888966e6d8f306078b095a5df801a583549d0/bitsandbytes-0.46.1-py3-none-manylinux_2_24_x86_64.whl", hash = "sha256:b0ee4a204fb926d4eae02bc2f5468ae3c11c011cfa849a4c771d4c6b201f57ae", size = 72865511, upload-time = "2025-07-02T19:44:25.192Z" }, + { url = "https://files.pythonhosted.org/packages/85/7d/06da01fac23a5032632dd7874b31c1d9b7b9af2314b2b07e5f99641950da/bitsandbytes-0.46.1-py3-none-win_amd64.whl", hash = "sha256:9f6f61376bd0e9780c5dc4ddee7d1f52cb10fe8034a1ea588611f4e8b87eb6a7", size = 72197544, upload-time = "2025-07-02T19:44:29.483Z" }, +] + [[package]] name = "blake3" version = "1.0.5" @@ -676,6 +690,7 @@ version = "0.0.1" source = { editable = "." } dependencies = [ { name = "accelerate" }, + { name = "bitsandbytes" }, { name = "datasets" }, { name = "deepspeed" }, { name = "einops" }, @@ -708,6 +723,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "accelerate", specifier = "==1.6.0" }, + { name = "bitsandbytes", specifier = ">=0.46.1" }, { name = "datasets", specifier = "==3.6.0" }, { name = "deepspeed", specifier = "==0.17.1" }, { name = "einops" },