From faac7dc35fe0f3def6d20dc4edba1b828149c8cf Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Fri, 27 Jun 2025 14:59:27 +0900
Subject: [PATCH] readme dataloading + comments cleanup

---
 README.md                               | 32 +++++++++++++++++++---
 scripts/short_ctx/gemma_qa_short_ctx.sh | 34 ++++++++++++++++++++++++
 src/ctx_to_lora/data/collator.py        |  7 +----
 src/ctx_to_lora/data/processing.py      | 35 +++++++++++--------------
 src/ctx_to_lora/eval_utils.py           |  2 --
 src/ctx_to_lora/trainer.py              |  4 +++
 6 files changed, 83 insertions(+), 31 deletions(-)
 create mode 100644 scripts/short_ctx/gemma_qa_short_ctx.sh

diff --git a/README.md b/README.md
index 28e28a9..504752c 100644
--- a/README.md
+++ b/README.md
@@ -42,7 +42,7 @@ WANDB_MODE=disabled run uv run intx_sft.py configs/fw_qa_v2_level_0_tiny.yaml --
 ```bash
 WANDB_MODE=disabled uv run intx_sft.py configs/squad.yaml --model_name_or_path=google/gemma-3-1b-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
@@ -58,7 +58,7 @@ WANDB_MODE=disabled uv run python intx_sft.py configs/self_gen_3_mini.yaml --mod
 ### HyperLoRA w/ fw-qa pretrain only
 ```bash
 WANDB_MODE=disabled run uv run python intx_sft.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
@@ -97,6 +97,28 @@ ds = load_and_process_dataset(f"self_gen/{base_model_name}/{ds_name}",
     streaming=False,
     num_proc=8
 )
+
+# or load the tokenized version
+base_model_max_len = 2**13
+ctx_model_max_len = 2**13
+
+# load via a custom function because of custom chat_template
+tokenizer = get_tokenizer(base_model_name)
+ctx_tokenizer = get_tokenizer(base_model_name)
+ds = def get_tokenized_dataset(
+    ds_name
+    split="train",
+    base_model_max_len=base_model_max_len,
+    tokenizer=tokenizer,
+    tokenizer_kwargs={},
+    ctx_model_max_len=ctx_model_max_len,
+    ctx_tokenizer=ctx_tokenizer,
+    ctx_tokenizer_kwargs={},
+    add_ctx_to_chat=False,
+    add_repeat_prompt=False,
+    add_negative_prompt=False,
+    use_kl_loss=False, 
+)
 ```
 
 ***Generate data from scratch***
@@ -118,8 +140,10 @@ run uv run python data/generate_fw_edu_qa_v2.py --shard_pattern "000_00000" --n_
 ```bash
 # Example commands using gemma-2-2b-it
 # self-gen data for fw_qa_v2
-# 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_v2/min_0_to_2000/013*'
+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/013*_level_3*'
+
+# 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
diff --git a/scripts/short_ctx/gemma_qa_short_ctx.sh b/scripts/short_ctx/gemma_qa_short_ctx.sh
new file mode 100644
index 0000000..f9be2b7
--- /dev/null
+++ b/scripts/short_ctx/gemma_qa_short_ctx.sh
@@ -0,0 +1,34 @@
+#!/bin/bash
+#SBATCH --job-name=ctxlora
+#SBATCH --nodes=1
+#SBATCH --partition=sakura-gpu
+#SBATCH --gpus=4
+#SBATCH --output=slurm_logs/%x-%j.out
+#SBATCH --error=slurm_logs/%x-%j.out
+
+uv run accelerate launch --num_processes=4 --gradient_accumulation_steps=2 --gradient_clipping=1.0 \
+--gpu_ids all --main_process_port 29571 intx_sft.py configs/qa_short_ctx.yaml \
+--model_name_or_path=google/gemma-2-2b-it \
+--num_train_epochs=5 \
+--per_device_train_batch_size=-1 \
+--gradient_accumulation_steps=2 \
+--per_device_eval_batch_size=64 \
+--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 \
+--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=50
\ No newline at end of file
diff --git a/src/ctx_to_lora/data/collator.py b/src/ctx_to_lora/data/collator.py
index f3cde31..ae6af55 100644
--- a/src/ctx_to_lora/data/collator.py
+++ b/src/ctx_to_lora/data/collator.py
@@ -64,7 +64,6 @@ def train_collator(inp_list, tokenizer):
     sample = inp_list[0]
     ctx_ids = None
     if "ctx_ids" in sample:
-        # have to be manual since it has [ctx_len, features] shape
         # pad to the longest ctx_len in the batch
         # which can have a different length from the input_ids, attn_mask, labels
 
@@ -74,7 +73,6 @@ def train_collator(inp_list, tokenizer):
             batch_first=True,
             padding_value=0,
         )
-        # exotic keys won't be padded, so we need to pad them as well
         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,
@@ -110,6 +108,7 @@ def train_collator(inp_list, tokenizer):
 
 
 def eval_collator(inp_list, tokenizer):
+    # only used for teacher-forcing eval
     # input is a list of tokenized sequences
     padding_kwargs = dict(
         padding=True,
@@ -120,7 +119,6 @@ def eval_collator(inp_list, tokenizer):
 
     ctx_ids = None
     if "ctx_ids" in inp_list[0]:
-        # have to be manual since it has [ctx_len, features] shape
         # pad to the longest ctx_len in the batch
         # which can have a different length from the input_ids, attn_mask, labels
 
@@ -130,7 +128,6 @@ def eval_collator(inp_list, tokenizer):
             batch_first=True,
             padding_value=0,
         )
-        # exotic keys won't be padded, so we need to pad them as well
         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,
@@ -200,7 +197,6 @@ def generation_collator(inp_list, tokenizer):
     # since we don't want to compute the loss on the labels
     # during generation
     if "ctx_ids" in inp_list[0]:
-        # have to be manual since it has [ctx_len, features] shape
         # 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]
@@ -209,7 +205,6 @@ def generation_collator(inp_list, tokenizer):
             batch_first=True,
             padding_value=0,
         )
-        # exotic keys won't be padded, so we need to pad them as well
         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,
diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py
index 752347c..7128139 100644
--- a/src/ctx_to_lora/data/processing.py
+++ b/src/ctx_to_lora/data/processing.py
@@ -747,26 +747,23 @@ def construct_and_tokenize_ctx_qa(
         )
 
     # for use_kl_loss, we need "chat_ids" and "chat_attn_mask"
+    # TODO (distillation): implement
     if use_kl_loss:
-        raise NotImplementedError("KL loss deprecated")
-        tokenized_ds = tokenized_ds.map(
-            convert_ctx_prompt_response_to_messages,
-            fn_kwargs={"add_ctx_to_chat": True},
-            remove_columns=["messages"],
-        )
-        tokenized_ds = tokenized_ds.map(
-            get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer),
-            remove_columns=["chat"],
-        )
-        tokenized_ds = tokenized_ds.map(
-            tokenize_chat_messages,
-            fn_kwargs={
-                "tokenizer": tokenizer,
-                "mask_assistant_inputs": True,
-                "tokenizer_kwargs": tokenizer_kwargs,
-                "for_kl_loss": True,
-            },
-        )
+        raise NotImplementedError("KL loss not implemented")
+
+        # old code from before (might not work anymore)
+        # e.g,. the next line likely replace the tokenized text from
+        # earlier `convert_ctx_prompt_response_to_messages` call
+
+        # tokenized_ds = tokenized_ds.map(
+        #     convert_ctx_prompt_response_to_messages,
+        #     fn_kwargs={"add_ctx_to_chat": True},
+        #     remove_columns=["messages"],
+        # )
+        # tokenized_ds = tokenized_ds.map(
+        #     get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer),
+        #     remove_columns=["chat"],
+        # )
 
     if need_ctx_ids:
         # tokenize the ctx_text to get ctx_ids and ctx_attn_mask
diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py
index 6599e5b..9214ce6 100644
--- a/src/ctx_to_lora/eval_utils.py
+++ b/src/ctx_to_lora/eval_utils.py
@@ -444,8 +444,6 @@ def decode_test_result(
     """Decode test results into human-readable format."""
     out = []
     for sample, pred_toks in zip(test_dataset, test_result.predictions):
-        # sample and labels are not padded
-        # pred_toks are padded though
         d = dict()
         if "labels" in sample:
             start_idx = np.argmax(sample["labels"] != -100)
diff --git a/src/ctx_to_lora/trainer.py b/src/ctx_to_lora/trainer.py
index d2a71f8..e25534e 100644
--- a/src/ctx_to_lora/trainer.py
+++ b/src/ctx_to_lora/trainer.py
@@ -10,6 +10,10 @@ from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
 logger = logging.getLogger()
 
 
+# TODO (distillation): implement
+class DistillationTrainer(Trainer): ...
+
+
 class ModulatedModelTrainer(Trainer):
     def __init__(self, *args, **kwargs):
         self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)