self-gen data misaligned fixed + niah reproducible + vllm-tokenizers version combo

This commit is contained in:
51616 2025-10-08 05:17:50 +00:00
parent 39ed1b3d75
commit a074fd3305
8 changed files with 64 additions and 62 deletions

View file

@ -53,11 +53,14 @@ print(tokenizer.decode(outputs[0]))
### 🎮 Interactive Demo [WIP]
```bash
uv run webui/demo.py
uv run demo/app.py
```
### 🧪 Experimental Scripts
| Experiment | Data prep | Training | Evaluation | Notes |
| ------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- |
| [Main experiment](scripts/main_exp/) | `uv run bash scripts/main_exp/0-download_data.sh`<br>`uv run bash scripts/main_exp/gen_data.sh` *(optional, rebuilds from scratch)* | `uv run bash scripts/main_exp/1-train.sh` | `uv run bash scripts/main_exp/eval/<script>.sh` (pick from `base_model.sh`, `cd.sh`, `cd_oracle.sh`, `d2l.sh`, `llmlingua.sh`, `t2l.sh`) | Downloading data is fastest; regenerate only if you need fresh synthetic data. Evaluation scripts reproduce the main paper metrics. |
| [NIAH](scripts/niah/) | `uv run bash scripts/niah/0-gen_data.sh` | `uv run bash scripts/niah/1-train.sh` | `uv run bash scripts/niah/2-eval.sh` | Run the scripts in order; data generation only needs to happen once |
To run any of the following scripts, use `uv run $PATH_TO_SCRIPT` from the root of this project.
| Experiment | Data prep | Training | Evaluation | Notes |
| ------------------------------------ | ------------------------------------- | ----------------------------- | ---------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- |
| [Main experiment](scripts/main_exp/) | `scripts/main_exp/0-download_data.sh` | `scripts/main_exp/1-train.sh` | `scripts/main_exp/eval/*.sh` | Downloading data is fastest; regenerate only if you need fresh synthetic data. Evaluation scripts reproduce the main paper metrics. |
| [NIAH](scripts/niah/) | `scripts/niah/0-gen_data.sh` | `scripts/niah/1-train.sh` | `scripts/niah/2-eval.sh` | Run the scripts in order; data generation only needs to happen once |

View file

@ -233,7 +233,7 @@ def self_generate(
tk = get_tokenizer(args.vllm_model, train=True)
self_qa_intx_tokens = tk(SELF_QA_INTX, add_special_tokens=False)["input_ids"]
self_qa_intx_tokens = tk(SELF_QA_INTX, add_special_tokens=False)["input_ids"][1:]
if args.remove_qa_template:
self_qa_intx_tokens = tk("\n\n", add_special_tokens=False)["input_ids"]
n_self_qa_intx_tokens = len(self_qa_intx_tokens)

View file

@ -7,6 +7,7 @@ uv self update
uv venv --python 3.10 --seed
uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --torch-backend=cu124
uv sync
uv pip install tokenizers==0.21.0
uv pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
uv pip install flashinfer-python==0.2.2 -i https://flashinfer.ai/whl/cu124/torch2.6

View file

@ -30,7 +30,7 @@ dependencies = [
"pandas",
"plotly",
"rouge-score",
"vllm==0.8.4",
"vllm==0.8.5.post1",
"huggingface-hub[hf-transfer]>=0.32.0",
"opt-einsum>=3.4.0",
"kagglehub[hf-datasets]>=0.3.12",
@ -38,7 +38,6 @@ dependencies = [
"bitsandbytes>=0.46.1",
"google-cloud-storage>=3.2.0",
"wonderwords>=2.2.0",
"tokenizers==0.21.0",
"llmlingua>=0.2.2",
]

View file

@ -1,22 +1,40 @@
WANDB_MODE=disabled uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=1 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=32 \
--per_device_eval_batch_size=16 \
--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 --logging_steps=10 \
--save_steps=1000 --learning_rate=4e-5 \
--lora_dropout=0.0 --neftune_noise_alpha=0 \
--per_rank_gen=True --per_layer_processing=True \
--gen_lora_l1_reg_coef=1 --use_sequence_packing=True \
--max_packed_inp_len=2048 --max_packed_ctx_len=2048 \
--dataloader_num_workers=0 --dataloader_prefetch_factor=None \
--ctx_encoder_type=early_exit --n_latent_queries=208 \
--eval_on_start=True --lora_r=8 \
--max_ctx_chunk_len=512 --min_ctx_chunk_len=25 \
--num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' \
--max_val_samples_per_ds=100 --seed=1 --use_per_ctx_average_loss=True
WANDB_MODE=disabled run uv run train.py \
configs/niah_exp/ctx_magic_number_32_256.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=1 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=16 \
--per_device_eval_batch_size=16 \
--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 \
--logging_steps=10 \
--save_steps=1000 \
--learning_rate=4e-5 \
--lora_dropout=0.0 \
--neftune_noise_alpha=0 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=1.5 \
--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=False \
--eval_on_start=True \
--max_ctx_chunk_len=512 \
--min_ctx_chunk_len=25 \
--num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' \
--max_val_samples_per_ds=100 \
--seed=1 \
--use_per_ctx_average_loss=True \
--torch_empty_cache_steps=10

1
scripts/niah/2-eval-test.sh Executable file
View file

@ -0,0 +1 @@
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path CHECKPOINT_PATH --datasets ctx_magic_number_32_1024 ctx_magic_number_1024_2048 ctx_magic_number_3072_4096 ctx_magic_number_7168_8192 ctx_magic_number_15360_16384 ctx_magic_number_28672_32768 ctx_magic_number_57344_65536 ctx_magic_number_122880_131072 --max_ctx_chunk_len=1024 --split test --eval_batch_size_gen=4

8
scripts/niah/README.md Normal file
View file

@ -0,0 +1,8 @@
# NIAH experiment
```bash
# run the scripts in this order
# data generation is only needed to be run once
uv run bash scripts/niah/0-gen_data.sh
uv run bash scripts/niah/1-train.sh
uv run bash scripts/niah/2-eval.sh
```

36
uv.lock generated
View file

@ -856,7 +856,6 @@ dependencies = [
{ name = "setuptools" },
{ name = "tensorboard" },
{ name = "tensorboardx" },
{ name = "tokenizers" },
{ name = "torchmetrics" },
{ name = "transformers" },
{ name = "vllm" },
@ -894,10 +893,9 @@ requires-dist = [
{ name = "setuptools" },
{ name = "tensorboard" },
{ name = "tensorboardx" },
{ name = "tokenizers", specifier = "==0.21.0" },
{ name = "torchmetrics" },
{ name = "transformers", specifier = "==4.51.3" },
{ name = "vllm", specifier = "==0.8.4" },
{ name = "vllm", specifier = "==0.8.5.post1" },
{ name = "wandb" },
{ name = "wonderwords", specifier = ">=2.2.0" },
]
@ -5428,30 +5426,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/e6/34/ebdc18bae6aa14fbee1a08b63c015c72b64868ff7dae68808ab500c492e2/tinycss2-1.4.0-py3-none-any.whl", hash = "sha256:3a49cf47b7675da0b15d0c6e1df8df4ebd96e9394bb905a5775adb0d884c5289", size = 26610, upload-time = "2024-10-24T14:58:28.029Z" },
]
[[package]]
name = "tokenizers"
version = "0.21.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "huggingface-hub" },
]
sdist = { url = "https://files.pythonhosted.org/packages/20/41/c2be10975ca37f6ec40d7abd7e98a5213bb04f284b869c1a24e6504fd94d/tokenizers-0.21.0.tar.gz", hash = "sha256:ee0894bf311b75b0c03079f33859ae4b2334d675d4e93f5a4132e1eae2834fe4", size = 343021, upload-time = "2024-11-27T13:11:23.89Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/b0/5c/8b09607b37e996dc47e70d6a7b6f4bdd4e4d5ab22fe49d7374565c7fefaf/tokenizers-0.21.0-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:3c4c93eae637e7d2aaae3d376f06085164e1660f89304c0ab2b1d08a406636b2", size = 2647461, upload-time = "2024-11-27T13:11:07.911Z" },
{ url = "https://files.pythonhosted.org/packages/22/7a/88e58bb297c22633ed1c9d16029316e5b5ac5ee44012164c2edede599a5e/tokenizers-0.21.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:f53ea537c925422a2e0e92a24cce96f6bc5046bbef24a1652a5edc8ba975f62e", size = 2563639, upload-time = "2024-11-27T13:11:05.908Z" },
{ url = "https://files.pythonhosted.org/packages/f7/14/83429177c19364df27d22bc096d4c2e431e0ba43e56c525434f1f9b0fd00/tokenizers-0.21.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6b177fb54c4702ef611de0c069d9169f0004233890e0c4c5bd5508ae05abf193", size = 2903304, upload-time = "2024-11-27T13:10:51.315Z" },
{ url = "https://files.pythonhosted.org/packages/7e/db/3433eab42347e0dc5452d8fcc8da03f638c9accffefe5a7c78146666964a/tokenizers-0.21.0-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:6b43779a269f4629bebb114e19c3fca0223296ae9fea8bb9a7a6c6fb0657ff8e", size = 2804378, upload-time = "2024-11-27T13:10:53.513Z" },
{ url = "https://files.pythonhosted.org/packages/57/8b/7da5e6f89736c2ade02816b4733983fca1c226b0c42980b1ae9dc8fcf5cc/tokenizers-0.21.0-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:9aeb255802be90acfd363626753fda0064a8df06031012fe7d52fd9a905eb00e", size = 3095488, upload-time = "2024-11-27T13:11:00.662Z" },
{ url = "https://files.pythonhosted.org/packages/4d/f6/5ed6711093dc2c04a4e03f6461798b12669bc5a17c8be7cce1240e0b5ce8/tokenizers-0.21.0-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d8b09dbeb7a8d73ee204a70f94fc06ea0f17dcf0844f16102b9f414f0b7463ba", size = 3121410, upload-time = "2024-11-27T13:10:55.674Z" },
{ url = "https://files.pythonhosted.org/packages/81/42/07600892d48950c5e80505b81411044a2d969368cdc0d929b1c847bf6697/tokenizers-0.21.0-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:400832c0904f77ce87c40f1a8a27493071282f785724ae62144324f171377273", size = 3388821, upload-time = "2024-11-27T13:10:58.401Z" },
{ url = "https://files.pythonhosted.org/packages/22/06/69d7ce374747edaf1695a4f61b83570d91cc8bbfc51ccfecf76f56ab4aac/tokenizers-0.21.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e84ca973b3a96894d1707e189c14a774b701596d579ffc7e69debfc036a61a04", size = 3008868, upload-time = "2024-11-27T13:11:03.734Z" },
{ url = "https://files.pythonhosted.org/packages/c8/69/54a0aee4d576045b49a0eb8bffdc495634309c823bf886042e6f46b80058/tokenizers-0.21.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:eb7202d231b273c34ec67767378cd04c767e967fda12d4a9e36208a34e2f137e", size = 8975831, upload-time = "2024-11-27T13:11:10.32Z" },
{ url = "https://files.pythonhosted.org/packages/f7/f3/b776061e4f3ebf2905ba1a25d90380aafd10c02d406437a8ba22d1724d76/tokenizers-0.21.0-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:089d56db6782a73a27fd8abf3ba21779f5b85d4a9f35e3b493c7bbcbbf0d539b", size = 8920746, upload-time = "2024-11-27T13:11:13.238Z" },
{ url = "https://files.pythonhosted.org/packages/d8/ee/ce83d5ec8b6844ad4c3ecfe3333d58ecc1adc61f0878b323a15355bcab24/tokenizers-0.21.0-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:c87ca3dc48b9b1222d984b6b7490355a6fdb411a2d810f6f05977258400ddb74", size = 9161814, upload-time = "2024-11-27T13:11:16.675Z" },
{ url = "https://files.pythonhosted.org/packages/18/07/3e88e65c0ed28fa93aa0c4d264988428eef3df2764c3126dc83e243cb36f/tokenizers-0.21.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:4145505a973116f91bc3ac45988a92e618a6f83eb458f49ea0790df94ee243ff", size = 9357138, upload-time = "2024-11-27T13:11:20.09Z" },
{ url = "https://files.pythonhosted.org/packages/15/b0/dc4572ca61555fc482ebc933f26cb407c6aceb3dc19c301c68184f8cad03/tokenizers-0.21.0-cp39-abi3-win32.whl", hash = "sha256:eb1702c2f27d25d9dd5b389cc1f2f51813e99f8ca30d9e25348db6585a97e24a", size = 2202266, upload-time = "2024-11-27T13:11:28.784Z" },
{ url = "https://files.pythonhosted.org/packages/44/69/d21eb253fa91622da25585d362a874fa4710be600f0ea9446d8d0217cec1/tokenizers-0.21.0-cp39-abi3-win_amd64.whl", hash = "sha256:87841da5a25a3a5f70c102de371db120f41873b854ba65e52bccd57df5a3780c", size = 2389192, upload-time = "2024-11-27T13:11:25.724Z" },
]
[[package]]
name = "tomli"
@ -5673,7 +5647,6 @@ dependencies = [
{ name = "regex" },
{ name = "requests" },
{ name = "safetensors" },
{ name = "tokenizers" },
{ name = "tqdm" },
]
sdist = { url = "https://files.pythonhosted.org/packages/f1/11/7414d5bc07690002ce4d7553602107bf969af85144bbd02830f9fb471236/transformers-4.51.3.tar.gz", hash = "sha256:e292fcab3990c6defe6328f0f7d2004283ca81a7a07b2de9a46d67fd81ea1409", size = 8941266, upload-time = "2025-04-14T08:15:00.485Z" }
@ -5877,7 +5850,7 @@ wheels = [
[[package]]
name = "vllm"
version = "0.8.4"
version = "0.8.5.post1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "aiohttp" },
@ -5925,7 +5898,6 @@ dependencies = [
{ name = "setuptools", marker = "python_full_version >= '3.12'" },
{ name = "six", marker = "python_full_version >= '3.12'" },
{ name = "tiktoken" },
{ name = "tokenizers" },
{ name = "torch" },
{ name = "torchaudio" },
{ name = "torchvision" },
@ -5936,9 +5908,9 @@ dependencies = [
{ name = "xformers", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "xgrammar", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/e6/d6/9d412cdaa92c3ab6250cef51217d37395b2aa372c6c14f90b1668adbbf63/vllm-0.8.4.tar.gz", hash = "sha256:522b13dd16c6c773dec0cb4c42ea591623d03ef94d16db8128ece2600017e6ac", size = 6667631, upload-time = "2025-04-15T00:30:10.147Z" }
sdist = { url = "https://files.pythonhosted.org/packages/68/a4/66914f97cd0b2e8b56d4ab9e1695415fbf38c3bc6f2dc5abf0487a32d2ef/vllm-0.8.5.post1.tar.gz", hash = "sha256:5e5be78ee00637de4ee29f75ce86edc6c224c05d9e58d067a511eb83c3afe32d", size = 7337621, upload-time = "2025-05-02T22:31:09.951Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/8e/cb/03dc1299e0456ff3d58a11f63682ef29aaf5b1bd7f21bfe0690d7ce6fc40/vllm-0.8.4-cp38-abi3-manylinux1_x86_64.whl", hash = "sha256:e346749ee8df48cdcd935d00a7fc123a1e17d9904b064401e74fc6ad73b8104a", size = 294098962, upload-time = "2025-04-15T00:30:02.293Z" },
{ url = "https://files.pythonhosted.org/packages/47/07/e8822ca825f88aec61a7af6210b85cd0e0a4d6f757081d3282f32b1b0912/vllm-0.8.5.post1-cp38-abi3-manylinux1_x86_64.whl", hash = "sha256:b7f0dbd46f82aac6b2e7489c2bbf1c90fdc2e56d3f1fff44549951db40814d5a", size = 326421557, upload-time = "2025-05-02T22:31:02.23Z" },
]
[[package]]