Add distillation training (#5)

- work with vllm self-gen data (everything is tokenized during self-gen)
- fix chat template!
- not use liger kernel for distillation (for both training and eval)

.gitignore

@@ -6,7 +6,10 @@ models/
 results/
 datasets/
 llm-comparator/
-/data/
+# Ignore data files but not scripts
+/data/processed_datasets/
+/data/raw_datasets/
+/data/distil/
 __pycache__/
 *egg-info/
 .vscode/
@@ -22,7 +25,6 @@ plots/
 *.safetensors
 *tfevents*
 watcher_state.yaml
-ds_config.json
 eval_results/
 .github/
 *.code-workspace

README.md

@@ -33,6 +33,11 @@ print(model.decode(outputs))
 ```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
 ```
+
+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
+ ```
 ### FWQA-v2 Level-0 Tiny
 ```bash
 WANDB_MODE=disabled run uv run intx_sft.py configs/fw_qa_v2_level_0_tiny.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=24000 --max_packed_ctx_len=48000 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=50
@@ -71,6 +76,12 @@ uv run python data/generate_fav_num.py
 uv run python data/generate_fav_num_big.py
 ```
 
+### SQuAD
+```bash
+# for some reason download directly through `load_dataset` does not work
+HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download --repo-type dataset rajpurkar/squad --local-dir data/raw_datasets/squad
+```
+
 ### Data Generation (v2)
 ***Download a subset of generated Fineweb-QA directly***
 ```bash

chat_templates/google/gemma-2-2b-it.jinja

@@ -14,7 +14,7 @@
     {%- if message['role'] == 'user' and loop.first and system_message is defined %}
         {{ '<start_of_turn>' + role + '\n' + system_message + '\n\n' + message['content'] | trim + '<end_of_turn>\n' }}
     {%- elif message['role'] == 'assistant' %}
-        {{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] | trim + '<end_of_turn>\n' }}{% endgeneration %}
+        {{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] + '<end_of_turn>' }}{% endgeneration %}{{ '\n' }}
     {%- else %}
         {{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}
     {%- endif %}

chat_templates/google/gemma-2-2b-it.jinja.bak(w_trim)

@@ -0,0 +1,24 @@
+{{- bos_token }}
+{%- if messages[0]['role'] == 'system' %}
+    {%- set system_message = messages[0]['content'] %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{% for message in loop_messages %}    
+    {% if (message['role'] == 'assistant') %}
+        {% set role = 'model' %}
+    {% else %}
+        {% set role = message['role'] %}
+    {% endif %}
+    {%- if message['role'] == 'user' and loop.first and system_message is defined %}
+        {{ '<start_of_turn>' + role + '\n' + system_message + '\n\n' + message['content'] | trim + '<end_of_turn>\n' }}
+    {%- elif message['role'] == 'assistant' %}
+        {{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] | trim + '<end_of_turn>\n' }}{% endgeneration %}
+    {%- else %}
+        {{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}
+    {%- endif %}
+{% endfor %}
+{% if add_generation_prompt %}
+    {{'<start_of_turn>model\n'}}
+{% endif %}

configs/context_numbers_10_self_gen.yaml

@@ -0,0 +1,43 @@
+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
+# save_strategy: "no"
+# # save_steps: 500
+# logging_strategy: "steps"
+# logging_steps: 100
+# use_liger_kernel: true
+# remove_unused_columns: false
+
+# needed to avoid OOM by compute the metrics batch by batch
+# w/o this the trainer stores logits of all sample in memory...
+# batch_eval_metrics: true
+
+per_device_train_batch_size: 64
+per_device_eval_batch_size: 128
+max_new_tokens: 64
+gen_per_device_eval_batch_size: 128
+max_val_samples_per_ds: 500
+# optim: schedule_free_adamw
+learning_rate: 0.0001
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 1
+weight_decay: 0.01
+
+# LoRA
+lora_r: 16
+lora_dropout: 0.0
+target_modules:
+  - down_proj
+
+# data
+train_ds_names:
+- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/context_numbers_2_10
+
+val_ds_names:
+- context_numbers_2_10

configs/pwc_tiny.yaml

@@ -0,0 +1,45 @@
+output_dir: "" # just a placeholder
+bf16: true
+model_name_or_path: google/gemma-3-1b-it
+label_names: ["labels"]
+# eval_on_start: True
+# eval_strategy: "steps"
+# eval_steps: 500
+# save_strategy: "no"
+# # save_steps: 500
+# logging_strategy: "steps"
+# logging_steps: 100
+# use_liger_kernel: true
+# remove_unused_columns: false
+
+# needed to avoid OOM by compute the metrics batch by batch
+# w/o this the trainer stores logits of all sample in memory...
+# batch_eval_metrics: true
+
+per_device_train_batch_size: 8
+per_device_eval_batch_size: 8
+max_val_samples_per_ds: 1000
+# optim: schedule_free_adamw
+
+learning_rate: 0.00004
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 1
+weight_decay: 0.01
+
+warmup_steps: 100
+
+dataloader_prefetch_factor: 16
+dataloader_num_workers: 8
+
+# LoRA
+lora_r: 8
+lora_dropout: 0.0
+target_modules:
+  - down_proj
+
+# data
+train_ds_names:
+  - pwc_tiny
+
+val_ds_names:
+  - pwc

configs/qa_short_ctx_self_gen_no_fw_qa.yaml

@@ -39,10 +39,10 @@ target_modules:
 
 # data
 train_ds_names:
-  - 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/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

configs/qa_short_ctx_self_gen_squad.yaml

@@ -0,0 +1,45 @@
+output_dir: "" # just a placeholder
+bf16: true
+model_name_or_path: google/gemma-3-1b-it
+label_names: ["labels"]
+# eval_on_start: True
+# eval_strategy: "steps"
+# eval_steps: 500
+# save_strategy: "no"
+# # save_steps: 500
+# logging_strategy: "steps"
+# logging_steps: 100
+# use_liger_kernel: true
+# remove_unused_columns: false
+
+# needed to avoid OOM by compute the metrics batch by batch
+# w/o this the trainer stores logits of all sample in memory...
+# batch_eval_metrics: true
+
+per_device_train_batch_size: 8
+per_device_eval_batch_size: 8
+max_val_samples_per_ds: 1000
+# optim: schedule_free_adamw
+
+learning_rate: 0.00004
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 1
+weight_decay: 0.01
+
+warmup_steps: 100
+
+dataloader_prefetch_factor: 16
+dataloader_num_workers: 8
+
+# LoRA
+lora_r: 8
+lora_dropout: 0.0
+target_modules:
+  - down_proj
+
+# data
+train_ds_names:
+  - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
+
+val_ds_names:
+  - squad

data/build_squad_compact.py

@@ -6,7 +6,7 @@ from tqdm import tqdm
 from ctx_to_lora.data.processing import closed_qa_prompting
 
 if __name__ == "__main__":
-    ds_name = "rajpurkar/squad"
+    ds_name = "data/raw_datasets/squad"
 
     for split in ["train", "validation"]:
         ctx_qa_dict = dict()

data/download_self_gen_qa_logprobs.py

@@ -0,0 +1,9 @@
+from huggingface_hub import snapshot_download
+
+if __name__ == "__main__":
+    fw_dir = "./data/raw_datasets/self_gen/"
+    snapshot_download(
+        "Rujikorn/self_gen_qa_logprobs",
+        repo_type="dataset",
+        local_dir=fw_dir,
+    )

data/prepare_distillation_dataset.py

@@ -0,0 +1,143 @@
+import argparse
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from ctx_to_lora.data.processing import load_and_process_dataset
+
+
+def add_logits(row, model, tok, device, args):
+    ctx = row["context"]
+    qs = row["prompts"]
+    ans = row["responses"]
+    responses_masks = []
+    all_logits_values = []
+    all_logits_positions = []
+
+    for question, answer in zip(qs, ans):
+        messages = [
+            {
+                "role": "user",
+                "content": f"Answer the question based on the following information:\n{ctx}.\n\n{question}",
+            },
+            {"role": "assistant", "content": answer},
+        ]
+        all_tokens = tok.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_special_tokens=False,
+            truncation=False,
+            padding=False,
+            add_generation_prompt=False,
+            return_dict=True,
+        )
+
+        input = torch.tensor([all_tokens["input_ids"]], device=device)
+        with torch.no_grad():
+            logits = model(input).logits
+
+        top_logits, top_logits_positions = torch.topk(
+            logits, args.max_logits_to_store, dim=-1
+        )
+
+        # Convert logits to specified precision
+        dtype_map = {
+            "float32": torch.float32,
+            "bfloat16": torch.bfloat16,
+            "float16": torch.float16,
+        }
+        target_dtype = dtype_map[args.logit_precision]
+        top_logits = top_logits.to(target_dtype)
+
+        tokens_to_mask = tok.apply_chat_template(
+            [messages[0]],  # Just the user message
+            tokenize=True,
+            add_special_tokens=False,
+            truncation=False,
+            padding=False,
+            add_generation_prompt=True,
+            return_dict=True,
+        )
+        number_of_tokens_to_mask = len(tokens_to_mask["input_ids"])
+
+        # Create mask: False for prompt tokens, True for response tokens
+        responses_mask = torch.zeros(len(all_tokens["input_ids"]), dtype=torch.bool)
+        responses_mask[number_of_tokens_to_mask:] = True
+
+        responses_masks.append(responses_mask)
+        all_logits_values.append(top_logits.cpu())
+        all_logits_positions.append(top_logits_positions.cpu())
+
+    row["labels_logits_values"] = all_logits_values
+    row["labels_logits_positions"] = all_logits_positions
+    row["responses_mask"] = responses_masks
+    return row
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ds_name", required=True, help="name of the dataset")
+    ap.add_argument("--split", required=True, help="split of the dataset")
+    ap.add_argument(
+        "--max_logits_to_store",
+        type=int,
+        default=100,
+        help="max number of logits kept per completion",
+    )
+    ap.add_argument(
+        "--limit_samples",
+        type=int,
+        default=None,
+        help="process only the first n dataset rows",
+    )
+    ap.add_argument(
+        "--logit_precision",
+        default="bfloat16",
+        choices=["float32", "bfloat16", "float16"],
+        help="dtype for saved logits",
+    )
+    ap.add_argument(
+        "--model_name", required=True, help="model name to use for generating logits"
+    )
+    ap.add_argument(
+        "--output_path", required=True, help="path to save the processed dataset"
+    )
+    args = ap.parse_args()
+
+    # Load model and tokenizer
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = AutoModelForCausalLM.from_pretrained(args.model_name)
+    model = model.to(device)
+    model.eval()
+    tok = AutoTokenizer.from_pretrained(args.model_name)
+
+    # Load dataset
+    ds = load_and_process_dataset(
+        args.ds_name,
+        args.split,
+        add_negative_prompt=False,
+        add_repeat_prompt=False,
+        repeat_prob=0,
+        is_pretrain=False,
+        streaming=False,
+        num_proc=8,
+    )
+
+    # Limit samples if specified
+    if args.limit_samples:
+        ds = ds.select(range(min(args.limit_samples, len(ds))))
+
+    # Process dataset with logits
+    ds = ds.map(
+        lambda row: add_logits(row, model, tok, device, args),
+        batched=False,
+        desc="Adding logits to dataset",
+    )
+
+    # Save the processed dataset
+    ds.save_to_disk(args.output_path)
+    print(f"Dataset saved to {args.output_path}")
+
+
+if __name__ == "__main__":
+    main()

data/self_generate_qa.py

@@ -4,7 +4,7 @@ import random
 import re
 from glob import glob
 
-import pandas as pd
+import numpy as np
 import yaml
 from datasets import Dataset, load_dataset
 from vllm import LLM, SamplingParams
@@ -13,12 +13,15 @@ from ctx_to_lora.data.definitions import (
     CLOSED_QA_INTX_TEMPLATES,
     RAW_DATA_DIR,
     SELF_GEN_DATA_DIR,
+    SELF_GEN_SYSTEM_MSG,
 )
 from ctx_to_lora.data.processing import (
     filter_none,
     get_preprocessing_fn,
     load_and_process_dataset,
+    tokenize_ctx_text,
 )
+from ctx_to_lora.model_loading import get_tokenizer
 from ctx_to_lora.utils import clear_gpu
 
 STOP_STRINGS = {
@@ -31,14 +34,55 @@ STOP_STRINGS = {
 SYSTEM_TEMPLATE = (
     "### SYSTEM INSTRUCTION ###\n"
     "You are an honest and helpful assistant.\n"
-    "You must use the information provided in the context for responding to the question.\n"
+    # "You must use the information provided in the context for responding to the question.\n"
     # "If the context does not contain enough information to answer the question, you must say so.\n"
     # "If that is the case, you can answer based on your knowledge, but you must clearly state that you are doing so.\n"
     "**DO NOT** hallucinate or make up information.\n"
     "### END OF SYSTEM INSTRUCTION ###"
 )
 
-PROMPT_TEMPLATE = "### Context ###\n{context}\n\n\n### Question ###\n{question}"
+SELF_QA_INTX = (
+    "---\n\n"
+    "# System Instruction\n"
+    "- The information provided is up-to-date information.\n"
+    "- When the provided information is not relevant to the question, ***ignore*** it and answer the question based on your knowledge.\n"
+    "- If the provided information is related to the question, incorporate it in your response.\n"
+    "\n---\n\n"
+    "# User\n"
+)
+
+PRE_CTX = "# Information\n"
+
+PROMPT_TEMPLATE = (
+    PRE_CTX
+    + "{context}\n\n---\n\n"
+    +
+    # "# System Instruction\n"
+    # "- The information provided is up-to-date information.\n"
+    # "- When the provided information is not relevant to the question, ***ignore*** it and answer the question based on your knowledge.\n"
+    # "- If the provided information is related to the question, incorporate it in your response.\n"
+    # "\n---\n\n"
+    # "# User\n{question}"
+    SELF_QA_INTX
+    + "{question}"
+)
+
+
+"""
+
+---
+
+# System Instruction
+- The information provided is up-to-date information.
+- When the provided information is not relevant to the question, ***ignore*** it and answer the question based on your knowledge.
+- If the provided information is related to the question, incorporate it in your response.
+
+---
+# User
+
+What team is Mbappe playing for?
+
+"""
 
 MODEL_CTX_LEN = {
     "google/gemma-2-27b-it": 8192,
@@ -64,12 +108,12 @@ def load_config(config_path: str) -> dict:
     return config
 
 
-def check_should_skip(txt: str, vllm_model: str) -> bool:
-    """Check if the response should be skipped based on stop strings."""
-    for stop in STOP_STRINGS[vllm_model]:
-        if stop in txt[-len(stop) :]:
-            return (txt.split(stop)[0], False)  # Found a valid stop string
-    return (txt, True)  # No valid stop string found, skip this response
+# def check_should_skip(txt: str, vllm_model: str) -> bool:
+#     """Check if the response should be skipped based on stop strings."""
+#     for stop in STOP_STRINGS[vllm_model]:
+#         if stop in txt[-len(stop) :]:
+#             return (txt.split(stop)[0], False)  # Found a valid stop string
+#     return (txt, True)  # No valid stop string found, skip this response
 
 
 def get_dataset_configs(
@@ -123,27 +167,27 @@ def create_messages(
     ctxs: list[str], questions: list[list[str]], vllm_model: str, system_template: str
 ) -> list[list[dict]]:
     """Create chat messages for the model."""
-    if "gemma" in vllm_model:
-        # gemma models do not support system messages
-        return [
-            [
-                {
-                    "role": "user",
-                    "content": system_template + "\n\n\n" + get_prompt(ctx, q),
-                }
-            ]
-            for ctx, q_list in zip(ctxs, questions)
-            for q in q_list
-        ]
-    else:
-        return [
-            [
-                {"role": "system", "content": system_template},
-                {"role": "user", "content": get_prompt(ctx, q)},
-            ]
-            for ctx, q_list in zip(ctxs, questions)
-            for q in q_list
+    # if "gemma" in vllm_model:
+    # gemma models do not support system messages
+    return [
+        [
+            {
+                "role": "user",
+                "content": system_template + "\n\n\n" + get_prompt(ctx, q),
+            }
         ]
+        for ctx, q_list in zip(ctxs, questions)
+        for q in q_list
+    ]
+    # else:
+    #     return [
+    #         [
+    #             {"role": "system", "content": system_template},
+    #             {"role": "user", "content": get_prompt(ctx, q)},
+    #         ]
+    #         for ctx, q_list in zip(ctxs, questions)
+    #         for q in q_list
+    #     ]
 
 
 def self_generate(
@@ -219,6 +263,25 @@ def self_generate(
 
     ds = ds.filter(filter_none, batched=False, num_proc=8)
 
+    tk = get_tokenizer(args.vllm_model, train=True)
+
+    self_qa_intx_tokens = tk(SELF_QA_INTX, add_special_tokens=False)["input_ids"]
+    n_self_qa_intx_tokens = len(self_qa_intx_tokens)
+    pre_ctx_tokens = tk(PRE_CTX, add_special_tokens=False)["input_ids"]
+    n_pre_ctx_tokens = len(pre_ctx_tokens)
+    sys_tokens = tk(system_template.split("\n")[0], add_special_tokens=False)[
+        "input_ids"
+    ]
+    n_sys_tokens = len(sys_tokens)
+    os.environ["TOKENIZERS_PARALLELISM"] = "true"
+    ds = ds.map(
+        tokenize_ctx_text,
+        fn_kwargs={"tokenizer": tk},
+        batched=True,
+        batch_size=50_000,
+        keep_in_memory=True,
+    )
+
     ctxs = [sample["context"] for sample in ds]
     questions = [
         [add_closed_qa_prompt(q, closed_qa_prob) for q in sample["prompts"] if q]
@@ -226,59 +289,212 @@ def self_generate(
     ]
     print(f"Loaded {len(ctxs)} contexts and {len(questions)} questions")
 
-    messages = create_messages(ctxs, questions, args.vllm_model, SYSTEM_TEMPLATE)
+    k = 16
+    fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}/{ds_name}/{split}/ds{shard_name}"
 
-    print(f"Generating from {len(messages)} contexts")
-    # TODO (distillation): make vllm outputs logits here too
+    chunk_size = 1_000
+    for chunk_idx, start in enumerate(range(0, len(ctxs), chunk_size)):
+        print(f"Processing chunk {chunk_idx}")
+
+        chunk_ctxs = ctxs[start : start + chunk_size]
+        chunk_questions = questions[start : start + chunk_size]
+        chunk_messages = create_messages(
+            chunk_ctxs, chunk_questions, args.vllm_model, SELF_GEN_SYSTEM_MSG
+        )
+        print(f"Generating from {len(chunk_messages)} contexts")
+
+        # Clear GPU memory before processing the next chunk
+        clear_gpu()
+        execute_qa_generation(
+            fpath + f"_{chunk_idx:04d}",
+            args,
+            llm,
+            temp,
+            tk,
+            self_qa_intx_tokens,
+            n_self_qa_intx_tokens,
+            sys_tokens,
+            n_sys_tokens,
+            chunk_ctxs,
+            ds["ctx_ids"][start : start + chunk_size],
+            chunk_questions,
+            chunk_messages,
+            k,
+        )
+
+
+def execute_qa_generation(
+    fpath,
+    args,
+    llm,
+    temp,
+    tk,
+    self_qa_intx_tokens,
+    n_self_qa_intx_tokens,
+    sys_tokens,
+    n_sys_tokens,
+    ctxs,
+    ctx_ids,
+    questions,
+    messages,
+    k,
+):
     completions = llm.chat(
         messages,
         sampling_params=SamplingParams(
             max_tokens=1024,
+            logprobs=k,
             temperature=temp,
-            # needed for checking if stop tokens are present
+            seed=42,
+            spaces_between_special_tokens=False,
             skip_special_tokens=False,
             include_stop_str_in_output=True,
         ),
     )
 
-    self_gen_data = {ctx: {"prompts": [], "responses": []} for ctx in ctxs}
+    self_gen_data = {
+        ctx: {
+            "ctx_ids": ctx_ids,
+            "input_ids": [],
+            "response_start_end": [],
+            "logprobs_vals": [],
+            "logprobs_indices": [],
+        }
+        for ctx, ctx_ids in zip(ctxs, ctx_ids)
+    }
     c = 0
     n_skips = 0
+    sys_start = None
     for ctx, q_list in zip(ctxs, questions):
-        for i, q in enumerate(q_list):
-            response = completions[c + i].outputs[0].text
-            response, skip = check_should_skip(response, args.vllm_model)
-            # skip = True
-            # for stop in STOP_STRINGS[args.vllm_model]:
-            #     if stop == response[-len(stop) :]:
-            #         # Check if response ends with stop string
-            #         response = response.split(stop)[0]
-            #         skip = False
-            #         break
-            if skip:
-                print(f"Skipping due to missing stop string")
+        # self_gen_data[ctx]["ctx_ids"] = ctx_ids
+        for i, _ in enumerate(q_list):
+            # response = completions[c + i].outputs[0].text
+            reason = completions[c + i].outputs[0].finish_reason
+            if reason != "stop":
+                # print(f"idx: {c + i}")
+                print(f"finish_reason: {completions[c + i].outputs[0].finish_reason}")
+                print(f"Skipping due to finish_reason={reason} != 'stop'")
                 n_skips += 1
                 continue
 
-            self_gen_data[ctx]["prompts"].append(q)
-            self_gen_data[ctx]["responses"].append(response)
+            # includes the logprob before the first response token
+            # but excludes the logprob from eos token
+            logp = completions[c + i].outputs[0].logprobs
+
+            # len = num response tokens
+            n_response_tokens = len(completions[c + i].outputs[0].token_ids)
+
+            logp_indices = np.empty((n_response_tokens, k), dtype=np.int32)
+            # float-16 is better for this range
+            logp_vals = np.empty((n_response_tokens, k), dtype=np.float16)
+            assert len(logp) == n_response_tokens, (
+                f"Expected {n_response_tokens} logp entries, got {len(logp)}"
+            )
+
+            for li, info_d in enumerate(logp):
+                for j, (idx, tok_info) in enumerate(info_d.items()):
+                    logp_indices[li, j] = idx
+                    logp_vals[li, j] = tok_info.logprob
+
+            prompt_ids = completions[c + i].prompt_token_ids  # 1d list
+            # token_ids only includes generated tokens, not the prompt
+            response_token_ids = completions[c + i].outputs[0].token_ids  # 1d list
+            all_ids = prompt_ids + response_token_ids
+            res_start = len(prompt_ids)
+            res_end = res_start + n_response_tokens
+
+            if sys_start is None:
+                for ii in range(len(prompt_ids) - n_sys_tokens):
+                    if prompt_ids[ii : ii + n_sys_tokens] == sys_tokens:
+                        # found the start of the system message
+                        sys_start = ii
+                        break
+
+            q_start = None
+            for ii in range(
+                len(prompt_ids) - n_self_qa_intx_tokens,
+                -1,
+                -1,
+            ):
+                if prompt_ids[ii : ii + n_self_qa_intx_tokens] == self_qa_intx_tokens:
+                    # found the start of the user input
+                    q_start = ii + n_self_qa_intx_tokens
+                    break
+
+            # bos + question + eos + start model turn + response + eos
+            input_ids = all_ids[:sys_start] + all_ids[q_start:res_end]
+
+            # TODO: save prompt logprobs with qa_start
+            # TODO: also save the prompt tokens for output alignment during training
+            # prompt_logp = completions[c + i].prompt_logprobs
+            # q_len = res_start - q_start
+            # prompt_logp_indices = np.empty((q_len, k), dtype=np.int32)
+            # # float-16 is better for this range
+            # prompt_logp_vals = np.empty((q_len, k), dtype=np.float16)
+            # for li, info_d in enumerate(prompt_logp[q_start:res_start]):
+            #     for j, (idx, tok_info) in enumerate(info_d.items()):
+            #         # vllm returns logprob for gt token first,
+            #         # which means that we might have k + 1 logprobs
+            #         if j >= k:
+            #             continue
+            #         prompt_logp_indices[li, j] = idx
+            #         prompt_logp_vals[li, j] = tok_info.logprob
+
+            # relative to the input_ids
+            res_start = res_start - q_start + sys_start
+            res_end = res_start + n_response_tokens
+
+            # arrays will be saved as nested lists of numbers
+
+            self_gen_data[ctx]["input_ids"].append(input_ids)
+            # assume single-turn chat
+            self_gen_data[ctx]["response_start_end"].append((res_start, res_end))
+            self_gen_data[ctx]["logprobs_vals"].append(logp_vals)
+            self_gen_data[ctx]["logprobs_indices"].append(logp_indices)
+            # NOTE: also have to shift back 1 for training
+            # self_gen_data[ctx]["prompt_start_end"].append((sys_start, res_start))
+            # self_gen_data[ctx]["prompt_logprobs_vals"].append(prompt_logp_vals)
+            # self_gen_data[ctx]["prompt_logprobs_indices"].append(prompt_logp_indices)
+
         c += i + 1
 
     print(f"Skipped {n_skips} responses due to missing stop strings")
     samples = [
         {
-            "context": ctx,
-            "prompts": q_list,
-            "responses": self_gen_data[ctx]["responses"],
+            # "context": ctx,
+            # "prompts": q_list,
+            # "responses": self_gen_data[ctx]["responses"],
+            "ctx_ids": self_gen_data[ctx]["ctx_ids"],
+            "input_ids": self_gen_data[ctx]["input_ids"],
+            "response_start_end": self_gen_data[ctx]["response_start_end"],
+            # "prompt_start_end": self_gen_data[ctx]["prompt_start_end"],
+            "logprobs_vals": self_gen_data[ctx]["logprobs_vals"],
+            "logprobs_indices": self_gen_data[ctx]["logprobs_indices"],
         }
         for ctx, q_list in zip(ctxs, questions)
     ]
 
     if args.debug:
         for sample in samples:
-            print(f"context={sample['context']}")
-            print(f"prompt={sample['prompts']}")
-            print(f"response={sample['responses']}")
+            print(f"context={tk.decode(sample['ctx_ids'])}")
+            print(f"QA={[tk.decode(ids) for ids in sample['input_ids']]}")
+
+            for input_ids, (start, end) in zip(
+                sample["input_ids"], sample["response_start_end"]
+            ):
+                print(f"start={start}, end={end}")
+                print(f"response={tk.decode(input_ids[start:end])}")
+
+            for input_ids, (start, end) in zip(
+                sample["input_ids"], sample["prompt_start_end"]
+            ):
+                print(f"start={start}, end={end}")
+                print(f"prompt={tk.decode(input_ids[start:end])}")
+
+            print(f"logprobs_vals={[x.shape for x in sample['logprobs_vals']]}")
+            print(f"logprobs_indices={[x.shape for x in sample['logprobs_indices']]}")
+            for indices in sample["logprobs_indices"]:
+                print(f"logprobs_indices={indices[-1]}")
             print("=" * 80)
 
     print(f"Generated {len(samples)} samples")
@@ -286,16 +502,18 @@ def self_generate(
 
     # Save results
     if not args.debug:
-        df = pd.DataFrame(samples)
-        ds_out = Dataset.from_pandas(df)
-        fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}/{ds_name}/{split}/ds{shard_name}"
+        # df = pd.DataFrame(samples)
+        # ds_out = Dataset.from_pandas(df)
+        ds_out = Dataset.from_list(samples)
+        # fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}/{ds_name}/{split}/ds{shard_name}"
+
         os.makedirs(os.path.dirname(fpath), exist_ok=True)
         fpath = f"{fpath}.parquet"
         ds_out.to_parquet(fpath)
         print(f"Saved to {fpath}")
 
         # Cleanup
-        del samples, df, ds_out, completions, messages, ctxs, questions
+        del samples, ds_out, completions, messages, ctxs, questions
         clear_gpu()
 
 
@@ -369,6 +587,8 @@ if __name__ == "__main__":
         enable_prefix_caching=True,
         enable_chunked_prefill=True,
         max_model_len=MODEL_CTX_LEN.get(vllm_model),
+        max_num_batched_tokens=16384,
+        max_num_seqs=32,  # avoid oom when getting logprobs
     )
 
     print(f"{llm_kwargs=}")

ds_config.json

@@ -0,0 +1,20 @@
+{
+    "bf16": {
+        "enabled": "auto"
+    },
+    "zero_optimization": {
+        "stage": 0,
+        "stage3_gather_16bit_weights_on_model_save": "no",
+        "offload_optimizer": {
+            "device": "auto"
+        },
+        "offload_param": {
+            "device": "auto"
+        }
+    },
+    "gradient_clipping": "auto",
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "gradient_accumulation_steps": "auto",
+    "steps_per_print": 2000000
+}

intx_sft.py

@@ -17,7 +17,6 @@ from transformers import (
     AutoConfig,
     set_seed,
 )
-from transformers.utils import is_liger_kernel_available
 
 from ctx_to_lora.configs import (
     AggregatorArguments,
@@ -45,6 +44,7 @@ from ctx_to_lora.metrics import (
 )
 from ctx_to_lora.model_loading import (
     get_lora_config,
+    get_model,
     get_model_and_tokenizer,
     get_tokenizer,
 )
@@ -176,7 +176,7 @@ def main():
     ############ Model setup
     if not ctx_args.from_pretrained_checkpoint:
         model_name = model_args.model_name_or_path
-        model, tokenizer = get_model_and_tokenizer(
+        base_model, tokenizer = get_model_and_tokenizer(
             **vars(model_args),
             train=True,
             requires_grad=False,  # ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE,
@@ -191,8 +191,8 @@ def main():
                 ctx_encoder_model_config = ctx_encoder_model_config.text_config
             ctx_tokenizer = get_tokenizer(ctx_name, train=True)
         else:
-            ctx_name = model.base_model.config.name_or_path
-            ctx_encoder_model_config = model.config
+            ctx_name = base_model.base_model.config.name_or_path
+            ctx_encoder_model_config = base_model.config
             ctx_tokenizer = tokenizer
 
     if ctx_args.exp_setup == ExperimentSetup.HYPERLORA:
@@ -200,7 +200,7 @@ def main():
         logger.info("Using HyperLoRA")
         if not ctx_args.from_pretrained_checkpoint:
             hypernet_config = get_hypernet_config(
-                model,
+                base_model,
                 ctx_encoder_model_config,
                 hypernet_args,
                 aggregator_args,
@@ -215,13 +215,13 @@ def main():
                     " as the context encoder"
                 )
             model = ModulatedPretrainedModel(
-                model,
+                base_model,
                 hypernet_config,
                 ctx_encoder_args,
                 ctx_args.use_kl_loss,
                 ctx_args.use_sequence_packing,
             )
-            training_args.gen_lora_l1_reg_coef = ctx_args.gen_lora_l1_reg_coef
+
         else:
             logger.info(
                 f"Loading from checkpoint: {ctx_args.from_pretrained_checkpoint}"
@@ -237,7 +237,9 @@ def main():
             if ctx_name is None:
                 ctx_name = model.base_model.config.name_or_path
             ctx_tokenizer = get_tokenizer(ctx_name, train=True)
-            training_args.gen_lora_l1_reg_coef = ctx_args.gen_lora_l1_reg_coef
+
+        training_args.gen_lora_l1_reg_coef = ctx_args.gen_lora_l1_reg_coef
+        training_args.use_kl_loss = ctx_args.use_kl_loss
 
         if len([p for p in model.ctx_encoder.parameters() if p.requires_grad]):
             raise ValueError("ctx_encoder contains trainable parameters")
@@ -268,6 +270,10 @@ def main():
     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,
         max_qas_len=ctx_args.max_qas_len,
@@ -284,6 +290,7 @@ def main():
         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()}
@@ -409,33 +416,24 @@ def main():
     # HACK: [local path]: make liger kernel works with gemma3
     # see .venv/lib/python3.10/site-packages/liger_kernel/transformers/monkey_patch.py
     # L810:814 + L781
-    if training_args.use_liger_kernel and is_liger_kernel_available():
-        from liger_kernel.transformers import _apply_liger_kernel_to_instance
 
-        if isinstance(model, ModulatedPretrainedModel):
-            logger.info("Applying liger-kernel to ModulatedPretrainedModel")
-            if isinstance(model.base_model, PeftModel):
-                base_model = model.base_model.base_model
-            else:
-                base_model = model.base_model
-
-            if ctx_name is not None:
-                logger.info("Applying liger-kernel to ctx_encoder_model")
-                ctx_base_model = model.ctx_encoder.base_model
-                _apply_liger_kernel_to_instance(model=ctx_base_model)
-                compile_linear(ctx_base_model)
-                # model.ctx_encoder.compile(fullgraph=True)
-
-        elif isinstance(model, PeftModel):
-            logger.info("Applying liger-kernel to PeftModel")
+    if isinstance(model, ModulatedPretrainedModel):
+        if isinstance(model.base_model, PeftModel):
+            base_model = model.base_model.base_model
+        else:
             base_model = model.base_model
 
-        _apply_liger_kernel_to_instance(model=base_model.model)
-        # compile_linear(base_model)
-        base_model.compile(fullgraph=True, mode="max-autotune")
-        base_model.loss_function = torch.compile(
-            base_model.loss_function, fullgraph=True, mode="max-autotune"
-        )
+        if ctx_name is not None:
+            logger.info("Compiling ctx_encoder_model")
+            ctx_base_model = model.ctx_encoder.base_model
+            compile_linear(ctx_base_model)
+            # model.ctx_encoder.compile(fullgraph=True, mode="max-autotune")
+
+    elif isinstance(model, PeftModel):
+        base_model = model.base_model
+
+    logger.info("Compiling base_model")
+    base_model.compile(fullgraph=True, mode="max-autotune")
 
     if LOCAL_RANK == 0:
         wandb.init(

scripts/short_ctx/gemma_qa_short_ctx_old_conf_distill.sh

@@ -0,0 +1,43 @@
+#!/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=$((10000 + ($SLURM_JOBID % 50000)))
+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 \
+--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 \
+--n_latent_queries=208 \
+--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 \
+--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=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

src/ctx_to_lora/configs.py

@@ -233,7 +233,7 @@ class TrainingArguments(TrainingArguments):
         default=100,
     )
     use_liger_kernel: bool = field(
-        default=True,
+        default=False,
     )
     remove_unused_columns: bool = field(
         default=False,

src/ctx_to_lora/data/collator.py

@@ -103,6 +103,14 @@ def train_collator(inp_list, tokenizer):
     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
 

src/ctx_to_lora/data/definitions.py

@@ -10,6 +10,16 @@ TRANSFORMED_DATA_DIR = "data/processed_datasets"
 RAW_DATA_DIR = "data/raw_datasets/"
 SELF_GEN_DATA_DIR = f"{RAW_DATA_DIR}/self_gen/"
 
+# see https://huggingface.co/datasets/YuxinJiang/LTE_train_data/viewer/default/train?row=97&views%5B%5D=train&sql_row=2
+# based on https://arxiv.org/pdf/2402.11905
+SELF_GEN_SYSTEM_MSG = (
+    "### SYSTEM INSTRUCTION ###\n"
+    "You are an honest and helpful assistant.\n"
+    "You must use the information provided in the context for responding to the question.\n"
+    "**DO NOT** hallucinate or make up information.\n"
+    "### END OF SYSTEM INSTRUCTION ###"
+)
+
 
 REPEAT_PROMPTS = [
     "Repeat the text.",
@@ -124,23 +134,22 @@ DS_KWARGS = {
             split="train[60:]",
         ),
     ),
-    # "pwc_tiny": dict(
-    #     train=dict(path="sggetao/PwC", split="train[900:2000]"),
-    #     validation=dict(path="sggetao/PwC", split="train[:900]"),
-    # ),
-    "squad": dict(
+    "pwc_compact_tiny": dict(
         train=dict(
-            path="/home/rujikorn_sakana_ai/.cache/huggingface/hub/datasets--rajpurkar--squad",
-            split="train[900:]",
-        ),
-        validation=dict(
-            path="/home/rujikorn_sakana_ai/.cache/huggingface/hub/datasets--rajpurkar--squad",
-            split="train[:900]",
-        ),
-        test=dict(
-            path="/home/rujikorn_sakana_ai/.cache/huggingface/hub/datasets--rajpurkar--squad",
-            split="validation",
+            path="parquet",
+            data_files="data/raw_datasets/pwc_compact/train/ds.parquet",
+            # correspond to skipping to first 900 samples in the original dataset
+            split="train[60:200]",
         ),
+    ),
+    "pwc_tiny": dict(
+        train=dict(path="sggetao/PwC", split="train[900:2000]"),
+        validation=dict(path="sggetao/PwC", split="train[:900]"),
+    ),
+    "squad": dict(
+        train=dict(path="data/raw_datasets/squad", split="train[900:]"),
+        validation=dict(path="data/raw_datasets/squad", split="train[:900]"),
+        test=dict(path="data/raw_datasets/squad", split="validation"),
         # train=dict(path="rajpurkar/squad", split="train[900:]"),
         # validation=dict(path="rajpurkar/squad", split="train[:900]"),
         # test=dict(path="rajpurkar/squad", split="validation"),

src/ctx_to_lora/data/packing.py

@@ -95,8 +95,26 @@ 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
+
+    if has_logits:
+        sequences = zip(
+            batch["input_ids"],
+            batch["labels"],
+            batch["logprobs_vals"],
+            batch["logprobs_indices"],
+        )
+        n_labels = sum(len(y) for x in batch["logprobs_vals"] for y in x)
+        k = len(batch["logprobs_vals"][0][0][0])  # assuming all have same k
+        logprobs_vals = np.empty((n_labels, k), dtype=np.float32)
+        logprobs_indices = np.empty((n_labels, k), dtype=np.int32)
+        logits_offset = 0
+    else:
+        sequences = zip(batch["input_ids"], batch["labels"])
+
     offset = ctx_offset = 0
-    for input_ids_b, labels_b in zip(batch["input_ids"], batch["labels"]):
+    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
@@ -120,6 +138,17 @@ def pack_data_points_FA(
         labels[inp_start:inp_end] = labels_b
         offset += inp_len
 
+        if has_logits:
+            logprobs_vals_b, logprobs_indices_b = sample[2:]
+            logprobs_vals_b = concat_list(logprobs_vals_b)
+            logprobs_indices_b = concat_list(logprobs_indices_b)
+            logits_len = len(logprobs_vals_b)
+            logprobs_vals[logits_offset : logits_offset + logits_len] = logprobs_vals_b
+            logprobs_indices[logits_offset : logits_offset + logits_len] = (
+                logprobs_indices_b
+            )
+            logits_offset += logits_len
+
     for ctx_ids_b in batch["ctx_ids"]:
         ctx_len = len(ctx_ids_b)
         ctx_start, ctx_end = ctx_offset, ctx_offset + ctx_len
@@ -127,13 +156,17 @@ def pack_data_points_FA(
         ctx_position_ids[ctx_start:ctx_end] = np.arange(ctx_len, dtype=np.long)
         ctx_offset += ctx_len
 
-    return {
+    out = {
         "ctx_ids": ctx_ids,
         "ctx_position_ids": ctx_position_ids,
         "input_ids": input_ids,
         "position_ids": position_ids,
         "labels": labels,
     }
+    if has_logits:
+        out["logprobs_vals"] = logprobs_vals
+        out["logprobs_indices"] = logprobs_indices
+    return out
 
 
 def pack_batch(
@@ -143,6 +176,8 @@ def pack_batch(
     max_packed_size: int = -1,
 ) -> dict[str, any]:
     need_flatten = check_is_iterable(batch["input_ids"][0][0])
+    assert need_flatten
+
     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"]]
@@ -179,6 +214,10 @@ def pack_batch(
         "labels": [],
         "n_queries": [],
     }
+    has_logits = "logprobs_vals" in batch
+    if has_logits:
+        packed_batch["logprobs_vals"] = []
+        packed_batch["logprobs_indices"] = []
 
     packing_efficiency_ratios = []
     ctx_packing_efficiency_ratios = []
@@ -192,6 +231,11 @@ def pack_batch(
             "input_ids": batch["input_ids"][start_idx:end_idx],
             "labels": batch["labels"][start_idx:end_idx],
         }
+        if has_logits:
+            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_batch["ctx_ids"].append(packed_item["ctx_ids"])
         packed_batch["ctx_position_ids"].append(packed_item["ctx_position_ids"])
@@ -199,6 +243,9 @@ def pack_batch(
         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:
+            packed_batch["logprobs_vals"].append(packed_item["logprobs_vals"])
+            packed_batch["logprobs_indices"].append(packed_item["logprobs_indices"])
 
         if max_packed_inp_len > 0:
             inp_efficiency = len(packed_item["input_ids"]) / max_packed_inp_len

src/ctx_to_lora/data/processing.py

@@ -9,8 +9,10 @@ from os import path
 from typing import Any
 
 import datasets
+import numpy as np
+import torch
 from datasets import Dataset, is_caching_enabled, load_dataset
-from transformers import PreTrainedTokenizerBase
+from transformers import PreTrainedModel, PreTrainedTokenizerBase
 
 from ctx_to_lora.data.definitions import (
     CLOSED_QA_INTX_TEMPLATES,
@@ -20,6 +22,7 @@ from ctx_to_lora.data.definitions import (
     RAW_DATA_DIR,
     REPEAT_PROMPTS,
     SELF_GEN_DATA_DIR,
+    SELF_GEN_SYSTEM_MSG,
     TRANSFORMED_DATA_DIR,
 )
 from ctx_to_lora.data.packing import pack_batch
@@ -130,7 +133,7 @@ def get_preprocessing_fn(
                 "response": sample["answer"],
             }
 
-    elif ds_name == "pwc":
+    elif ds_name == "pwc" or ds_name == "pwc_tiny":
         # original pwc
         def f(sample):
             return {
@@ -309,6 +312,22 @@ def get_preprocessing_fn(
 
         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
 
 
@@ -551,7 +570,19 @@ def load_and_process_dataset(
     cols_to_remove = [
         col
         for col in ds.column_names
-        if col not in ["context", "prompts", "responses", "qas", "variation"]
+        if col
+        not in [
+            "context",
+            "prompts",
+            "responses",
+            "qas",
+            "variation",
+            "logprobs_vals",
+            "logprobs_indices",
+            "input_ids",
+            "ctx_ids",
+            "response_start_end",
+        ]
     ]
     is_eval = split != "train"
     ds = ds.map(
@@ -613,9 +644,9 @@ def get_tokenized_dataset(
     repeat_prob: float = 0.0,  # only used when add_repeat_prompt is True
     set_format: str | None = None,
     streaming: bool = False,
+    base_model: PreTrainedModel | None = None,
 ) -> dict[str, Any]:
-    # TODO (distillation, Tan): make this works with pre-computed logits
-    assert not use_kl_loss, "KL loss is deprecated"
+    # assert not use_kl_loss, "KL loss is deprecated"  # TO BE REMOVED
     if add_repeat_prompt:
         assert repeat_prob > 0, f"add_repeat_prompt is set but repeat_prob = 0"
     logger.info(f"Loading dataset {ds_name} with split {split}...")
@@ -673,6 +704,7 @@ def get_tokenized_dataset(
         ds=ds,
         tokenizer=tokenizer,
         ctx_tokenizer=ctx_tokenizer,
+        model=base_model,
         num_proc=num_proc,
         **tokenize_kwargs,
     )
@@ -696,6 +728,7 @@ def construct_and_tokenize_ctx_qa(
     ctx_tokenizer_kwargs,
     add_ctx_to_chat,
     use_kl_loss,
+    model,
     need_ctx_ids,
     is_pretrain,
     is_paraphrased,
@@ -740,74 +773,104 @@ def construct_and_tokenize_ctx_qa(
 
         # for sft + chat_model, we need to convert the dataset to chat format
         # add "data" field
-        ds = ds.map(
-            convert_ctx_prompt_response_to_messages,
-            fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat},
-            num_proc=num_proc,
-            remove_columns=[col for col in ds.column_names if col != "context"],
-        )
-        # add input_ids, attention_mask, labels to "data"
-        os.environ["TOKENIZERS_PARALLELISM"] = "true"
-        logging.debug("Tokenizing inputs")
-        tokenized_ds = ds.map(
-            get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer),
-            batched=True,
-            batch_size=100_000,
-            # # num_proc=num_proc,
-        )
+        if "input_ids" in ds.column_names and "response_start_end" in ds.column_names:
+            # already tokenized dataset (e.g., self-gen qa data)
+            tokenized_ds = ds.map(
+                get_labels_from_input_ids,
+                num_proc=num_proc,
+            )
+        else:
+            ds = ds.map(
+                convert_ctx_prompt_response_to_messages,
+                fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat},
+                num_proc=num_proc,
+                # remove_columns=[col for col in ds.column_names if col != "context"],
+            )
+            # add input_ids, attention_mask, labels to "data"
+            os.environ["TOKENIZERS_PARALLELISM"] = "true"
+            logging.debug("Tokenizing inputs")
+            tokenized_ds = ds.map(
+                get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer),
+                batched=True,
+                batch_size=100_000,
+                # # num_proc=num_proc,
+            )
 
-    if "train" in split:
-        # base model cant handle these samples naturally
-        # should skip
-        # tokenized_ds = tokenized_ds.filter(
-        #     len_filter,
-        #     fn_kwargs={
-        #         "max_length": base_model_max_len,
-        #         "keys": ["input_ids", "labels"],
-        #     },
-        #     # batched=False,
-        #     # batch_size=0,
-        #     num_proc=num_proc,
-        # )
-        tokenized_ds = tokenized_ds.map(
-            len_filter_data,
-            fn_kwargs={
-                "max_length": base_model_max_len,
-                "keys": ["input_ids"],
-            },
-            num_proc=num_proc,
-        )
+    # if "train" in split:
+    #     # base model cant handle these samples naturally
+    #     # should skip
+    #     # tokenized_ds = tokenized_ds.filter(
+    #     #     len_filter,
+    #     #     fn_kwargs={
+    #     #         "max_length": base_model_max_len,
+    #     #         "keys": ["input_ids", "labels"],
+    #     #     },
+    #     #     # batched=False,
+    #     #     # batch_size=0,
+    #     #     num_proc=num_proc,
+    #     # )
+    #     tokenized_ds = tokenized_ds.map(
+    #         len_filter_data,
+    #         fn_kwargs={
+    #             "max_length": base_model_max_len,
+    #             "keys": ["input_ids"],
+    #         },
+    #         num_proc=num_proc,
+    #     )
 
-    # for use_kl_loss, we need "chat_ids" and "chat_attn_mask"
-    # TODO (distillation): implement
+    logger.info(tokenized_ds)
+
+    if use_kl_loss and ("train" in split):
+        if "logprobs_vals" not in ds.column_names:
+            assert model is not None, (
+                f"Base model is required for computing the "
+                "logits when `use_kl_loss` is True and logprobs are not included."
+            )
+            # add "data_logp"
+            tokenized_ds = tokenized_ds.map(
+                convert_ctx_prompt_response_to_messages,
+                fn_kwargs={"add_ctx_to_chat": True, "for_logp": True},
+                num_proc=num_proc,
+            )
+            # add logits info to "data_logp" field
+            logger.info("computing logprobs")
+            tokenized_ds = tokenized_ds.map(
+                get_sft_prompt_formatting_fn(
+                    TRAINING_TASK.COMPLETION, tokenizer, use_kl_loss, model
+                ),
+                batched=True,
+                batch_size=100_000,
+            )
+        else:
+            tokenized_ds = tokenized_ds.map(add_ctx_attn_mask, num_proc=num_proc)
+
+    keep_cols = [
+        "input_ids",
+        "attention_mask",
+        "labels",
+        "ctx_ids",
+        "ctx_attn_mask",
+        "context",
+        "data",
+        "data_logp",
+    ]
     if use_kl_loss:
-        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"],
-        # )
+        keep_cols += ["logprobs_vals", "logprobs_indices"]
+    tokenized_ds = tokenized_ds.remove_columns(
+        [col for col in tokenized_ds.column_names if col not in keep_cols],
+    )
 
     if need_ctx_ids:
-        # tokenize the ctx_text to get ctx_ids and ctx_attn_mask
-        os.environ["TOKENIZERS_PARALLELISM"] = "true"
-        logging.debug("Tokenizing context")
-        tokenized_ds = tokenized_ds.map(
-            tokenize_ctx_text,
-            fn_kwargs={"tokenizer": ctx_tokenizer},
-            batched=True,
-            batch_size=100_000,
-            # # num_proc=num_proc,
-        )
+        if "ctx_ids" not in tokenized_ds.column_names:
+            # tokenize the ctx_text to get ctx_ids and ctx_attn_mask
+            os.environ["TOKENIZERS_PARALLELISM"] = "true"
+            logging.debug("Tokenizing context")
+            tokenized_ds = tokenized_ds.map(
+                tokenize_ctx_text,
+                fn_kwargs={"tokenizer": ctx_tokenizer},
+                batched=True,
+                batch_size=100_000,
+            )
         if "train" in split:
             # TODO: do something if ctx length is longer than the ctx model length
             # e.g., drop or split for multi-lora training
@@ -823,8 +886,13 @@ def construct_and_tokenize_ctx_qa(
             tokenized_ds = tokenized_ds.map(
                 unpack_data,
                 num_proc=num_proc,
-                remove_columns=["data", "context"],
+                remove_columns=[
+                    col
+                    for col in tokenized_ds.column_names
+                    if col in ["data", "data_logp", "context"]
+                ],
             )
+            print(tokenized_ds)
 
             logging.debug(f"Split too long QAs with max length {max_qas_len}")
             tokenized_ds = tokenized_ds.map(
@@ -900,9 +968,57 @@ def build_intx_pretrain(sample: dict[str, str]):
     return {"text": sample["context"] + sample["qas"]}
 
 
+def add_ctx_attn_mask(sample: dict[str, Any]) -> dict[str, Any]:
+    """
+    Add context attention mask to the sample.
+
+    Args:
+        sample: A dictionary containing 'ctx_ids'
+
+    Returns:
+        The sample with 'ctx_attn_mask' added
+    """
+    if "ctx_ids" in sample:
+        sample["ctx_attn_mask"] = [1] * len(sample["ctx_ids"])
+    return sample
+
+
+def get_labels_from_input_ids(sample: dict[str, Any]) -> dict[str, Any]:
+    """
+    Extract labels from input_ids and response_start.
+
+    Args:
+        sample: A dictionary containing 'input_ids' and 'response_start'
+
+    Returns:
+        A dictionary with 'labels' field added
+    """
+    labels = []
+    attention_mask = []
+    for input_ids_i, (start_i, end_i) in zip(
+        sample["input_ids"], sample["response_start_end"]
+    ):
+        len_input_ids = len(input_ids_i)
+        # pad labels with -100
+        pad_len_left = start_i
+        pad_len_right = len_input_ids - end_i
+        labels.append(
+            [IGNORE_INDEX] * pad_len_left
+            + input_ids_i[start_i:]
+            + [IGNORE_INDEX] * pad_len_right
+        )
+        attention_mask.append([1] * len(input_ids_i))
+
+    sample["labels"] = labels
+    sample["attention_mask"] = attention_mask
+    return sample
+
+
 def get_sft_prompt_formatting_fn(
     sft_mode: TRAINING_TASK,
     tokenizer: PreTrainedTokenizerBase,
+    for_logp: bool = False,
+    model: PreTrainedModel | None = None,
 ) -> Callable[[dict[str, Any]], dict[str, Any]]:
     """
     Get a function that formats examples for supervised fine-tuning.
@@ -929,10 +1045,13 @@ def get_sft_prompt_formatting_fn(
             "Training with pre-training data is not supported yet."
         )
 
+    data_field = "data_logp" if for_logp else "data"
+
+    @torch.inference_mode()
     def f_intx(samples):
         # flatten all the messages into a list
         # tokenize, the pack back correctly
-        messages_list = [x for x in samples["data"]]
+        messages_list = [x for x in samples[data_field]]
 
         n_queries = [len(x) for x in messages_list]
         messages = concat_list(messages_list)
@@ -940,31 +1059,76 @@ def get_sft_prompt_formatting_fn(
         tokens = tokenizer.apply_chat_template(
             messages,
             tokenize=True,
-            add_special_token=False,
+            add_special_tokens=False,
+            padding=False,
             truncation=False,
             add_generation_prompt=False,
             return_assistant_tokens_mask=True,
             return_dict=True,
         )
         labels = []
+        logprobs_vals = []
+        logprobs_indices = []
         for tok_ids, masks in zip(tokens["input_ids"], tokens["assistant_masks"]):
             o = [id_ if mask else IGNORE_INDEX for id_, mask in zip(tok_ids, masks)]
             labels.append(o)
+            if for_logp:
+                # TODO: batch compute?
+                # save only the logits for the assistant tokens
+                assert sum(masks) > 0, "No assistant tokens found in the masks"
+
+                # + 1 to also get the logits of the token before the first label token
+                num_logits_to_keep = int(len(masks) - np.argmax(masks) + 1)
+
+                # [n_labels + 1]
+                logits = model(
+                    torch.tensor([tok_ids], device=model.device),
+                    logits_to_keep=num_logits_to_keep,
+                ).logits.squeeze(0)
+
+                # include the logits of the token before the first label token
+                pre_label_idx = np.argmax(masks) - 1
+                masks[pre_label_idx] = True
+                logits = logits[
+                    torch.tensor(
+                        masks[pre_label_idx:],
+                        device=logits.device,
+                        dtype=torch.bool,
+                    )
+                ]
+
+                logits = logits[..., :-1, :]  # remove the last token (eos)
+                logprobs = torch.nn.functional.log_softmax(logits.float(), dim=-1)
+                # ignore eos
+                vals, idx = torch.topk(logprobs, k=16, dim=-1)
+                logprobs_vals.append(vals.to(torch.float16))
+                logprobs_indices.append(idx.to(torch.int32))
+
         del tokens["assistant_masks"]
         tokens["labels"] = labels
         per_ctx_tokens = []
         i = 0
         for n in n_queries:
-            per_ctx_tokens.append(
-                {
+            if for_logp:
+                per_ctx_data = {
+                    "logprobs_vals": logprobs_vals[i : i + n],
+                    "logprobs_indices": logprobs_indices[i : i + n],
+                }
+                # per_ctx_data = {
+                #     "logprobs_vals": logprobs_vals[i : i + n],
+                #     "logprobs_indices": logprobs_indices[i : i + n],
+                # }
+            else:
+                per_ctx_data = {
                     "input_ids": tokens["input_ids"][i : i + n],
                     "attention_mask": tokens["attention_mask"][i : i + n],
                     "labels": tokens["labels"][i : i + n],
                 }
-            )
+
+            per_ctx_tokens.append(per_ctx_data)
             i += n
 
-        return dict(data=per_ctx_tokens)
+        return {data_field: per_ctx_tokens}
 
     # return f if not apply_chat_template_fn is not None else f_intx
     return f_intx
@@ -973,6 +1137,7 @@ def get_sft_prompt_formatting_fn(
 def convert_ctx_prompt_response_to_messages(
     example: dict[str, Any],
     add_ctx_to_chat: bool,
+    for_logp: bool = False,
 ) -> dict[str, Any]:
     """
     Convert context/prompt/response format to chat messages format.
@@ -998,6 +1163,8 @@ def convert_ctx_prompt_response_to_messages(
     system_msg = ""
     if "system_message" in example:
         system_msg = example["system_message"].strip()
+        if for_logp and not system_msg:
+            system_msg = SELF_GEN_SYSTEM_MSG
 
     data = []
     for prompt, response in zip(example[prompt_field], example[res_field]):
@@ -1009,11 +1176,12 @@ def convert_ctx_prompt_response_to_messages(
             [
                 {"role": "system", "content": system_msg.strip()},
                 {"role": "user", "content": user_msg.strip()},
-                {"role": "assistant", "content": response.strip()},
+                {"role": "assistant", "content": response},
             ]
         )
+    out_field = "data_logp" if for_logp else "data"
 
-    return dict(data=data)
+    return {out_field: data}
 
 
 def split_too_long_qas(
@@ -1031,6 +1199,8 @@ def split_too_long_qas(
     labels = samples["labels"]
     ctx_ids = samples["ctx_ids"]
     ctx_attn_mask = samples["ctx_attn_mask"]
+    target_logprobs_vals = samples.get("logprobs_vals", None)
+    target_logprobs_indices = samples.get("logprobs_indices", None)
 
     # Pre-calculate total lengths to check if any splitting is needed
     total_lengths = [sum(len(x) for x in seq) for seq in input_ids]
@@ -1048,14 +1218,28 @@ def split_too_long_qas(
     out = {k: list() for k in samples}
     longest_new_qas_len = 0
     n_skip = 0
+    has_target_logprobs = (
+        target_logprobs_vals is not None and target_logprobs_indices is not None
+    )
 
     # Helper function to add a batch efficiently
-    def add_batch(inp_ids_batch, attn_batch, labels_batch, ctx_id, ctx_attn):
+    def add_batch(
+        inp_ids_batch,
+        attn_batch,
+        labels_batch,
+        ctx_id,
+        ctx_attn,
+        target_vals_batch=None,
+        target_indices_batch=None,
+    ):
         out["input_ids"].append(inp_ids_batch)
         out["attention_mask"].append(attn_batch)
         out["labels"].append(labels_batch)
         out["ctx_ids"].append(ctx_id)
         out["ctx_attn_mask"].append(ctx_attn)
+        if has_target_logprobs:
+            out["logprobs_vals"].append(target_vals_batch)
+            out["logprobs_indices"].append(target_indices_batch)
 
     for i, tot_inp_len in enumerate(total_lengths):
         if (max_qas_len < 0 or tot_inp_len <= max_qas_len) and (
@@ -1074,10 +1258,26 @@ def split_too_long_qas(
         new_input_ids = []
         new_attn_mask = []
         new_labels = []
+        new_target_vals = [] if has_target_logprobs else None
+        new_target_indices = [] if has_target_logprobs else None
+
+        sequences = zip(input_ids[i], attention_mask[i], labels[i])
+        if has_target_logprobs:
+            sequences = zip(
+                input_ids[i],
+                attention_mask[i],
+                labels[i],
+                target_logprobs_vals[i],
+                target_logprobs_indices[i],
+            )
+
+        for seq_data in sequences:
+            if has_target_logprobs:
+                inp_ids, attn_mask, label, target_vals, target_indices = seq_data
+            else:
+                inp_ids, attn_mask, label = seq_data
+                target_vals, target_indices = None, None
 
-        for inp_ids, attn_mask, label in zip(
-            input_ids[i], attention_mask[i], labels[i]
-        ):
             inp_len = len(inp_ids)
             if inp_len > max_qas_len:
                 # Skip individual sequences that are too long
@@ -1095,6 +1295,9 @@ def split_too_long_qas(
                 new_input_ids.append(inp_ids)
                 new_attn_mask.append(attn_mask)
                 new_labels.append(label)
+                if has_target_logprobs:
+                    new_target_vals.append(target_vals)
+                    new_target_indices.append(target_indices)
             else:
                 # Current batch is full, save it and start new batch
                 # logger.debug(
@@ -1107,6 +1310,8 @@ def split_too_long_qas(
                         new_labels,
                         current_ctx_id,
                         current_ctx_attn,
+                        new_target_vals,
+                        new_target_indices,
                     )
                     longest_new_qas_len = max(longest_new_qas_len, new_qas_len)
 
@@ -1115,6 +1320,9 @@ def split_too_long_qas(
                 new_input_ids = [inp_ids]
                 new_attn_mask = [attn_mask]
                 new_labels = [label]
+                if has_target_logprobs:
+                    new_target_vals = [target_vals]
+                    new_target_indices = [target_indices]
 
         # Add final batch if not empty
         if new_input_ids:
@@ -1124,6 +1332,8 @@ def split_too_long_qas(
                 new_labels,
                 current_ctx_id,
                 current_ctx_attn,
+                new_target_vals,
+                new_target_indices,
             )
             longest_new_qas_len = max(longest_new_qas_len, new_qas_len)
 
@@ -1138,7 +1348,12 @@ def split_too_long_qas(
 
 
 def unpack_data(sample):
-    return {k: sample["data"][k] for k in sample["data"]}
+    if "data" not in sample:
+        return sample
+    data_keys = ["data"]
+    if "data_logp" in sample:
+        data_keys.append("data_logp")
+    return {k: sample[dk][k] for dk in data_keys for k in sample[dk]}
 
 
 def unpack_data_eval(samples):
@@ -1317,6 +1532,7 @@ def tokenize_pretrain(
     tokens = tokenizer(
         samples["text"],
         add_special_tokens=True,
+        padding=False,
         truncation=False,
         **(tokenizer_kwargs or {}),
     )
@@ -1438,6 +1654,72 @@ def pack(
     return ds
 
 
+# def add_logits(model, ds):
+#     # compute logits for prepared tokenized ds
+#     # must handle packed and flattened dataset
+#     # also cache!
+#     # let's do this row by row
+#     # slower for non-packed ds
+#     input = torch.tensor([ds['input_ids']], device=model.device)
+#     labels = torch.tensor([ds['labels']], device=model.device)
+#     with torch.no_grad():
+#         logits = model(input).logits
+
+#     top_logits, top_logits_positions = torch.topk(logits, args.max_logits_to_store, dim=-1)
+
+#     # Convert logits to specified precision
+#     dtype_map = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}
+#     target_dtype = dtype_map[args.logit_precision]
+#     top_logits = top_logits.to(target_dtype)
+
+#     tokens_to_mask = tok.apply_chat_template(
+#         [messages[0]],  # Just the user message
+#         tokenize=True,
+#         add_special_tokens=False,
+#         truncation=False,
+#         padding=False,
+#         add_generation_prompt=True,
+#         return_dict=True,
+#     )
+#     number_of_tokens_to_mask = len(tokens_to_mask['input_ids'])
+
+#     # Create mask: False for prompt tokens, True for response tokens
+#     responses_mask = torch.zeros(len(all_tokens['input_ids']), dtype=torch.bool)
+#     responses_mask[number_of_tokens_to_mask:] = True
+
+#     responses_masks.append(responses_mask)
+#     all_logits_values.append(top_logits.cpu())
+#     all_logits_positions.append(top_logits_positions.cpu())
+
+#     logits = model(input).logits
+
+#     top_logits, top_logits_positions = torch.topk(logits, args.max_logits_to_store, dim=-1)
+
+#     # Convert logits to specified precision
+#     dtype_map = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}
+#     target_dtype = dtype_map[args.logit_precision]
+#     top_logits = top_logits.to(target_dtype)
+
+#     tokens_to_mask = tok.apply_chat_template(
+#         [messages[0]],  # Just the user message
+#         tokenize=True,
+#         add_special_tokens=False,
+#         truncation=False,
+#         padding=False,
+#         add_generation_prompt=True,
+#         return_dict=True,
+#     )
+#     number_of_tokens_to_mask = len(tokens_to_mask['input_ids'])
+
+#     # Create mask: False for prompt tokens, True for response tokens
+#     responses_mask = torch.zeros(len(all_tokens['input_ids']), dtype=torch.bool)
+#     responses_mask[number_of_tokens_to_mask:] = True
+
+#     responses_masks.append(responses_mask)
+#     all_logits_values.append(top_logits.cpu())
+#     all_logits_positions.append(top_logits_positions.cpu())
+
+
 if __name__ == "__main__":
     from transformers import AutoTokenizer
 

src/ctx_to_lora/eval_utils.py

@@ -14,15 +14,12 @@ import pandas as pd
 import torch
 import yaml
 from datasets import disable_caching
-from peft import PeftModel
 from transformers import (
-    PreTrainedModel,
     Seq2SeqTrainer,
     Seq2SeqTrainingArguments,
     Trainer,
     set_seed,
 )
-from transformers.utils import is_liger_kernel_available
 
 from ctx_to_lora.data.collator import eval_collator, generation_collator
 from ctx_to_lora.data.definitions import (
@@ -693,24 +690,24 @@ def evaluate(
             use_flash_attn=True,
         )
 
-    if is_liger_kernel_available():
-        from liger_kernel.transformers import _apply_liger_kernel_to_instance
+    # if is_liger_kernel_available():
+    #     from liger_kernel.transformers import _apply_liger_kernel_to_instance
 
-        if isinstance(model, ModulatedPretrainedModel):
-            print("Applying liger-kernel to ModulatedPretrainedModel")
-            if isinstance(model.base_model, PeftModel):
-                _apply_liger_kernel_to_instance(model=model.base_model.base_model.model)
-            else:
-                _apply_liger_kernel_to_instance(model=model.base_model.model)
-            if ctx_name is not None:
-                print("Applying liger-kernel to ctx_encoder_model")
-                _apply_liger_kernel_to_instance(model=model.ctx_encoder.base_model)
-        elif isinstance(model, PeftModel):
-            print("Applying liger-kernel to PeftModel")
-            _apply_liger_kernel_to_instance(model=model.base_model.model)
-        elif isinstance(model, PreTrainedModel):
-            print("Applying liger-kernel to PretrainedModel")
-            _apply_liger_kernel_to_instance(model=model.model)
+    #     if isinstance(model, ModulatedPretrainedModel):
+    #         print("Applying liger-kernel to ModulatedPretrainedModel")
+    #         if isinstance(model.base_model, PeftModel):
+    #             _apply_liger_kernel_to_instance(model=model.base_model.base_model.model)
+    #         else:
+    #             _apply_liger_kernel_to_instance(model=model.base_model.model)
+    #         if ctx_name is not None:
+    #             print("Applying liger-kernel to ctx_encoder_model")
+    #             _apply_liger_kernel_to_instance(model=model.ctx_encoder.base_model)
+    #     elif isinstance(model, PeftModel):
+    #         print("Applying liger-kernel to PeftModel")
+    #         _apply_liger_kernel_to_instance(model=model.base_model.model)
+    #     elif isinstance(model, PreTrainedModel):
+    #         print("Applying liger-kernel to PretrainedModel")
+    #         _apply_liger_kernel_to_instance(model=model.model)
 
     tokenizer = get_tokenizer(args.model_name_or_path, train=False)
     if tokenizer.pad_token_id is None:

src/ctx_to_lora/modeling/lora_layer.py

@@ -29,11 +29,12 @@ def lora_forward(
     B = B.repeat_interleave(n_qs, dim=0, output_size=tot_q)
 
     base_out = torch.nn.Linear.forward(self, x, *args, **kwargs)
+    x = x.to(A.dtype)
     delta_x = F.dropout(x, p=lora_dropout_p, training=self.training)
     delta_x = einsum(A, delta_x, "tot_q r d_in, tot_q s_len d_in -> tot_q s_len r")
     delta_x = einsum(B, delta_x, "tot_q d_out r, tot_q s_len r -> tot_q s_len d_out")
     delta_x = delta_x * scaling
-    return base_out + delta_x
+    return (base_out + delta_x).to(base_out.dtype)
 
 
 def lora_forward_packed(
@@ -52,7 +53,7 @@ def lora_forward_packed(
 ) -> Float[Tensor, "1 tot_len d_out"]:
     # bs of x should be 1 in this case
     base_out = torch.nn.Linear.forward(self, x, *args, **kwargs)
-
+    x = x.to(A.dtype)
     delta_x = F.dropout(x, p=lora_dropout_p, training=self.training)
     # # [tot_seq_len, r, d_in]
     # repeated_A = A.repeat_interleave(seq_lens, dim=0, output_size=seq_lens.sum())
@@ -73,7 +74,7 @@ def lora_forward_packed(
     )
     delta_x = delta_x * scaling
 
-    return base_out + delta_x
+    return (base_out + delta_x).to(base_out.dtype)
 
 
 def apply_lora_to_layers(

src/ctx_to_lora/trainer.py

@@ -1,5 +1,6 @@
 import logging
 
+import torch
 from torch import nn
 from transformers import Trainer
 from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
@@ -12,8 +13,126 @@ from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
 logger = logging.getLogger()
 
 
-# TODO (distillation, Shin): implement
-class DistillationTrainer(Trainer): ...
+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)
+        self.distillation_temperature = kwargs.pop("distillation_temperature", 1.0)
+        super().__init__(*args, **kwargs)
+
+    def compute_loss(
+        self, model, inputs, return_outputs=False, num_items_in_batch=None
+    ):
+        if self.model_accepts_loss_kwargs:
+            loss_kwargs = {}
+            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)
+        # 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
+        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]=}"
+        )
+
+        # 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"
+
+        # 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
+        #####
+
+        ##### unpack gen lora dict and compute regularization loss
+        l1_norm = 0
+        for module, lora in gen_loras.items():
+            l1_norm += lora["A"].abs().mean() + lora["B"].abs().mean()
+        # rough estimate of the losses (we only log the values from one step)
+        if (self.state.global_step == 1 and self.args.logging_first_step) or (
+            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 += self.gen_lora_l1_reg_coef * l1_norm
+        #####
+
+        if (
+            self.args.average_tokens_across_devices
+            and self.model_accepts_loss_kwargs
+            and num_items_in_batch is not None
+        ):
+            loss *= self.accelerator.num_processes
+
+        return (loss, outputs) if return_outputs else loss
 
 
 class ModulatedModelTrainer(Trainer):
@@ -40,6 +159,9 @@ class ModulatedModelTrainer(Trainer):
             if num_items_in_batch is not None:
                 loss_kwargs["num_items_in_batch"] = num_items_in_batch
             inputs = {**inputs, **loss_kwargs}
+        # 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)
         # Save past state if it exists
         # TODO: this needs to be fixed and made cleaner later.
@@ -86,7 +208,11 @@ class ModulatedModelTrainer(Trainer):
         loss += self.gen_lora_l1_reg_coef * l1_norm
         #####
 
-        if self.args.average_tokens_across_devices and self.model_accepts_loss_kwargs:
+        if (
+            self.args.average_tokens_across_devices
+            and self.model_accepts_loss_kwargs
+            and num_items_in_batch is not None
+        ):
             loss *= self.accelerator.num_processes
 
         return (loss, outputs) if return_outputs else loss
@@ -122,8 +248,6 @@ def train_model(
         checkpoint = training_args.resume_from_checkpoint
         logger.info(f"Resuming from the checkpoint: {checkpoint}")
 
-    is_modulated_model = isinstance(model, ModulatedPretrainedModel)
-    trainer_cls = Trainer if not is_modulated_model else ModulatedModelTrainer
     trainer_kwargs = dict(
         model=model,
         args=training_args,
@@ -132,10 +256,20 @@ def train_model(
         data_collator=train_collator,
         compute_metrics=compute_metrics,
     )
+
+    is_modulated_model = isinstance(model, ModulatedPretrainedModel)
+    trainer_cls = Trainer
     if is_modulated_model:
-        logger.info("Training with modulated model. Using CustomTrainer.")
+        logger.info("Training with modulated model. Using ModulatedModelTrainer.")
+        trainer_cls = ModulatedModelTrainer
         trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef
         del training_args.gen_lora_l1_reg_coef
+
+        if training_args.use_kl_loss:
+            logger.info("Training with distillation loss. Using DistillationTrainer.")
+            trainer_cls = DistillationTrainer
+            del training_args.use_kl_loss
+
     if training_args.auto_find_batch_size:
         # set the batch size to some high number
         # which will be lowered by the Trainer