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

per-rank bias + skip interleave if cached+ add bnb

Browse source
This commit is contained in:
51616 2025-08-09 15:20:37 +00:00
parent d6c649e13c
commit 178e7165ea
12 changed files with 302 additions and 160 deletions
Download patch file Download diff file Expand all files Collapse all files

32
README.md
View file

@ -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
```
<!--
### HyperLoRA w/ self-gen 3 mini
```bash
# self-gen sft
# 3_mini config
uv run python data/self_generate_qa.py \
--vllm_model=google/gemma-2-2b-it --config=configs/self_gen_3_mini.yaml
WANDB_MODE=disabled uv run python train.py configs/self_gen_3_mini.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=10 --per_device_train_batch_size=32 --gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=2 --lora_r=8 --eval_steps=5000 --save_steps=5000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1
```
### HyperLoRA w/ fw-qa pretrain only
```bash
WANDB_MODE=disabled run uv run python train.py configs/fw_qa_pretrain_only.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=4 --gradient_accumulation_steps=8 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=4 --num_latent_factor=1 --lora_r=8 --eval_steps=100 --save_steps=100 --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 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1
``` -->
<!-- ### HyperLoRA w/ context_numbers_128
```bash
WANDB_MODE=disabled run python hyperlora/train.py configs/context_numbers_128.yaml --model_name_or_path=meta-llama/Llama-3.2-1B-Instruct --num_train_epochs=10 --per_device_train_batch_size=64 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj
``` -->
### 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

42
configs/medium_exp/qa_short_ctx_self_gen_lv1_closed_1_and_lv3_l2l.yaml Normal file
View file

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

14
configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_l2l.yaml
View file

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

1
pyproject.toml
View file

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

29
scripts/short_ctx/gemma_qa_short_ctx_exp_default_8_gpus.sh Normal file
View file

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

8
src/ctx_to_lora/configs.py
View file

@ -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)."},

35
src/ctx_to_lora/data/processing.py
View file

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

18
src/ctx_to_lora/model_loading.py
View file

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

2
src/ctx_to_lora/modeling/aggregator.py
View file

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

218
src/ctx_to_lora/modeling/hypernet.py
View file

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

47
train.py
View file

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

16
uv.lock generated
View file

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