context distillation eval + timer

README.md

@@ -292,4 +292,4 @@ This bug causes the self-generated data to have two `<bos>` tokens (only the `in
 
 NOTE: previously we were using `vllm==0.8.5` which does not have the double `<bos>` problem.
 However it has another bug that randomly causes some detokenization error (see https://github.com/vllm-project/vllm/issues/17448)
-Therefore we, not knowing about the double `<bos>` bug, degraded `vllm` to `0.8.4`.
+Therefore we, not knowing about the double `<bos>` bug, recently downgraded `vllm` to `0.8.4`.

run_eval.py

@@ -77,6 +77,28 @@ if __name__ == "__main__":
         action="store_true",
         help="Remove context when evaluating the base model.",
     )
+    parser.add_argument(
+        "--use_cd",
+        action="store_true",
+        help="Use context distillation model for evaluation.",
+    )
+    parser.add_argument(
+        "--cd_update_iterations",
+        type=int,
+        default=20,
+        help="Number of update iterations for context distillation during evaluation",
+    )
+    parser.add_argument(
+        "--cd_use_gen_q",
+        action="store_true",
+        help="Use generated queries for context distillation training.",
+    )
+    parser.add_argument(
+        "--num_gen_q",
+        type=int,
+        default=None,
+        help="Number of generated queries to use for context distillation.",
+    )
 
     cli_args = vars(parser.parse_args())
     # setup_logging(output_dir, debug=os.getenv("DEBUG", False))

src/ctx_to_lora/configs.py

@@ -456,6 +456,10 @@ class DataArguments:
         default=None,
         metadata={"help": "Test dataset names."},
     )
+    max_train_samples_per_ds: int | None = field(
+        default=None,
+        metadata={"help": "Maximum number of training samples per dataset."},
+    )
     max_val_samples_per_ds: int | None = field(
         default=1000,
         metadata={"help": "Maximum number of validation samples per dataset."},

src/ctx_to_lora/data/collator.py

@@ -125,7 +125,6 @@ def generation_collator(inp_list, tokenizer):
     )
 
     if "ctx_ids" in inp_list[0]:
-        # TODO: handle chunked ctx_ids
         # 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]

src/ctx_to_lora/data/definitions.py

@@ -10,16 +10,6 @@ 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 ###"
-)
-
 
 SUMMARIZATION_PROMPTS = [
     "Summarize the article."

src/ctx_to_lora/data/processing.py

@@ -25,6 +25,7 @@ from ctx_to_lora.data.definitions import (
 )
 from ctx_to_lora.data.packing import pack_batch
 from ctx_to_lora.data.preprocessing_fn import get_preprocessing_fn
+from ctx_to_lora.data.self_gen_template import QA_PROMPT_TEMPLATE
 from ctx_to_lora.utils import check_is_iterable, concat_list
 
 logger = logging.getLogger()
@@ -215,6 +216,7 @@ def get_tokenized_dataset(
     add_ctx_to_chat: bool,
     use_kl_loss: bool,
     max_new_tokens: int = 256,
+    add_self_distill_template: bool = False,
     set_format: str | None = None,
 ) -> dict[str, Any]:
     if max_qas_len > 0:
@@ -280,6 +282,7 @@ def get_tokenized_dataset(
         ds=ds,
         tokenizer=tokenizer,
         ctx_tokenizer=ctx_tokenizer,
+        add_self_distill_template=add_self_distill_template,
         num_proc=num_proc,
         **tokenize_kwargs,
     )
@@ -313,6 +316,7 @@ def construct_and_tokenize_ctx_qa(
     ds,
     split,
     max_new_tokens,
+    add_self_distill_template=False,
     set_format=None,
     num_proc=None,
 ):
@@ -326,14 +330,24 @@ def construct_and_tokenize_ctx_qa(
         # add "messages_list" field
         ds = ds.map(
             convert_ctx_prompt_response_to_messages,
-            fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat},
+            fn_kwargs={
+                "add_ctx_to_chat": add_ctx_to_chat,
+                "add_self_distill_template": add_self_distill_template,
+            },
             num_proc=16,
         )
         # add `input_ids`, `attention_mask`, `labels`
         os.environ["TOKENIZERS_PARALLELISM"] = "true"
         logging.debug("Tokenizing inputs")
+        # HACK: check magic num niah dataset
+        prepend_bos = True  # see known issues in readme
+        sample = ds[0]
+        inp_txt = sample["prompts"][0]
+        if inp_txt == "What is the special magic number? Reply with only the number.":
+            logger.info("Detected Niah dataset, not prepending BOS")
+            prepend_bos = False
         tokenized_ds = ds.map(
-            get_sft_prompt_formatting_fn(tokenizer),
+            get_sft_prompt_formatting_fn(tokenizer, prepend_bos),
             batched=True,
             batch_size=100_000,
         )
@@ -384,11 +398,16 @@ def construct_and_tokenize_ctx_qa(
         }
         logging.info(f"Chunking context with {split_ctx_kwargs=}")
         tokenized_ds = tokenized_ds.map(
-            split_too_long_ctx, fn_kwargs=split_ctx_kwargs, num_proc=16
+            split_too_long_ctx,
+            fn_kwargs=split_ctx_kwargs,
+            num_proc=16,
         )
+
+        logging.info(f"Num samples after ctx chunking: {len(tokenized_ds)}")
         logging.info(
-            f"Avg. num chunks per ctx: {np.mean(list(map(len, tokenized_ds['ctx_ids'])))}"
+            f"Avg. num chunks per ctx: {np.mean(tokenized_ds['n_ctx_chunks'])}"
         )
+        tokenized_ds = tokenized_ds.remove_columns(["n_ctx_chunks"])
 
         split_qa_kwargs = {
             "max_qas_len": max_qas_len,
@@ -469,6 +488,7 @@ def get_labels_from_input_ids(sample: dict[str, Any]) -> dict[str, Any]:
 
 def get_sft_prompt_formatting_fn(
     tokenizer: PreTrainedTokenizerBase,
+    prepend_bos: bool,
 ) -> Callable[[dict[str, Any]], dict[str, Any]]:
     """
     Get a function that formats examples for supervised fine-tuning.
@@ -508,10 +528,13 @@ def get_sft_prompt_formatting_fn(
             return_dict=True,
         )
 
-        if tokenizer.name_or_path == "google/gemma-2-2b-it":
+        if prepend_bos:
             # HACK: add bos at the beginning (see `Known Issues` in README.md)
-            for tok_ids in tokens["input_ids"]:
-                tok_ids = [tokenizer.bos_token_id] + tok_ids
+            for i, (tok_ids, assistant_masks) in enumerate(
+                zip(tokens["input_ids"], tokens["assistant_masks"])
+            ):
+                tokens["input_ids"][i] = [tokenizer.bos_token_id] + tok_ids
+                tokens["assistant_masks"][i] = [0] + assistant_masks
 
         labels = []
         for tok_ids, masks in zip(tokens["input_ids"], tokens["assistant_masks"]):
@@ -535,6 +558,7 @@ def get_sft_prompt_formatting_fn(
 def convert_ctx_prompt_response_to_messages(
     example: dict[str, Any],
     add_ctx_to_chat: bool,
+    add_self_distill_template: bool = False,
 ) -> dict[str, Any]:
     """
     Convert context/prompt/response format to chat messages format.
@@ -565,7 +589,12 @@ def convert_ctx_prompt_response_to_messages(
     for prompt, response in zip(example[prompt_field], example[res_field]):
         user_msg = prompt.strip()
         if add_ctx_to_chat:
-            user_msg = example["context"].strip() + "\n\n" + user_msg
+            if add_self_distill_template:
+                user_msg = QA_PROMPT_TEMPLATE.format(
+                    context=example["context"].strip(), question=user_msg
+                )
+            else:
+                user_msg = example["context"].strip() + "\n\n" + user_msg
 
         messages_list.append(
             [
@@ -603,7 +632,7 @@ def split_too_long_ctx(
     ctx_ids = sample["ctx_ids"]
     # Early exits
     if chunk_len <= 0 and max_num_split is None:
-        return {"ctx_ids": [ctx_ids]}
+        return {"ctx_ids": [ctx_ids], "n_ctx_chunks": 1}
 
     n_chunks = None  # will be sampled (train) or derived (eval)
     if is_train and num_chunk_probs is not None:
@@ -632,7 +661,7 @@ def split_too_long_ctx(
     # Safety: at least 1
     n_chunks = max(1, n_chunks)
     if n_chunks == 1:
-        return {"ctx_ids": [ctx_ids]}
+        return {"ctx_ids": [ctx_ids], "n_ctx_chunks": 1}
 
     avg_len = ceil(len(ctx_ids) / n_chunks)
     chunks = [ctx_ids[i : i + avg_len] for i in range(0, len(ctx_ids), avg_len)]
@@ -646,7 +675,7 @@ def split_too_long_ctx(
         chunks[i] = prefix + chunks[i] + suffix
     chunks[-1] = prefix + chunks[-1]
 
-    return {"ctx_ids": chunks}
+    return {"ctx_ids": chunks, "n_ctx_chunks": len(chunks)}
 
 
 def split_too_long_qas(
@@ -934,7 +963,10 @@ def pack(
     else:
         train_ds = interleave_datasets(
             list(ds_dict.values()),
-            probabilities=get_ds_prob(train_ds_lens, total_samples),
+            # probabilities=get_ds_prob(
+            #     train_ds_lens, total_samples
+            # ),
+            probabilities=[l / total_samples for l in train_ds_lens],
             seed=seed,
             stopping_strategy="all_exhausted",
         )

src/ctx_to_lora/data/q_generation_prompts.py

@@ -0,0 +1,114 @@
+import random
+
+SYSTEM_PROMPT = (
+    "You are a creative and helpful assistant. "
+    "You will be given a context and you need to generate questions from the given context. "
+    "**DO NOT** hallucinate or make up information."
+)
+
+PROMPT_TEMPLATE = (
+    "### Context ###\n{context}\n\n"
+    "### Instruction ###\n{instruction}\n\n"
+    "### Rules ###\n"
+    "Phrases like 'based on the provided context', 'according to the context', etc., must not to appear in your response.\n\n"
+    # "2. The questions should not overlap. They should be diverse, covering many aspects of the context.\n"
+    # "3. Do not give away too much information in the questions. For example, ask 'Who is X?' instead of 'Who is X that did Y?' when Y is clear from the context.\n"
+    # "4. Ignore the text formatting of the context, e.g., bold, italic, underline, etc.\n"
+    # "5. Ignore typos, spacing, and grammatical errors in the context.\n"
+    # "6. Always use proper grammar and punctuation.\n"
+    # "7. Try to use different question forms and styles.\n"
+    "### Output Format ###\n"
+    "The question/instruction/task/request should be in the following format:\n\n"
+    "Message: {{question}}\n\n"
+    "Use this output template regardless of the type of the output."
+)
+
+# taken from https://github.com/HazyResearch/cartridges/blob/2ac563d79c2f3367a9e780a7bb0b3cf3039a8d50/cartridges/data/resources.py#L195
+# def get_structuring_seed_prompt() -> str:
+#     DATA_FORMATS = [
+#         "JSON",
+#         "YAML",
+#         "TOML",
+#         "INI",
+#         "XML",
+#         "plain text",
+#     ]
+
+#     data_format = random.choice(DATA_FORMATS)
+
+#     EXAMPLES = [
+#         dedent(f"""
+#             Can you structure the information in the context in the following format: {data_format}? Be sure to include precise information like any dates, times, names, and numerical values.
+#         """).strip(),
+#     ]
+
+#     example = random.choice(EXAMPLES)
+
+#     return dedent(f"""
+#         Please generate a single chat message instructing an LLM to structure the information in {data_format}. The message can follow the following template, filling in details from the context:
+
+#         '{example}'
+#     """).strip()
+
+
+# def get_summarization_seed_prompt() -> str:
+#     prompts = [
+#         dedent("""
+#             Please generate a single chat message instructing an LLM to summarize part of the context.
+#             Make sure the instruction is very explicit about the section of the context that you want to summarize.
+#             Include details (ids, names, titles, dates, etc.) that make it clear what you are asking about.
+#         """).strip(),
+#         dedent("""
+#             Please generate a single chat message instructing an LLM to summarize a section.
+#             Make sure the instruction is explicit about the section that should be summarized and the document it is from.
+#         """).strip(),
+#     ]
+#     prompt = random.choice(prompts)
+#     return prompt
+
+
+def get_question_seed_prompt() -> str:
+    prompts = [
+        (
+            "Generate a question for an LLM that will test its knowledge of the information in the context above. "
+            "In your question be sure to include details (ids, names, titles, dates, etc.) that make it clear what you are asking about. "
+            "Output only a single question. Do NOT include any other text or explanation other than the question."
+        ),
+        (
+            "Generate a message for an LLM that will test its knowledge of the information in the context above. "
+            "Be sure to include details (ids, names, titles, dates, etc.) in the question so that it can be answered without access to the context (i.e. closed-book setting). "
+            "Output only a single question. Do NOT include any other text or explanation other than the question."
+        ),
+        (
+            "You are helping to quiz a user about the information in the context. "
+            "Please generate a question about the subsection of the context above. "
+            "Be sure to include details (ids, names, titles, dates, etc.) in the question to make it clear what you are asking about. "
+            "Answer only with the question, do not include any other text."
+        ),
+    ]
+    prompt = random.choice(prompts)
+    return prompt
+
+
+def get_use_case_seed_prompt() -> str:
+    prompt = (
+        "Your primary goal is to think about practical, real-world tasks or applications that someone could achieve using the knowledge contained within the provided context. "
+        "Consider how a user might want to apply this information in a real-world scenario, not just recall it. Put yourself into someone else's shoes and think how you might use the information. "
+        "After considering potential use cases, your task will be to generate an instruction or task that reflects one of these downstream applications. "
+        "This instruction or task should be something a user, who has access to this context, might ask when trying to accomplish their specific goal. "
+        "Output only a single instruction or task. Do NOT include any other text or explanation."
+    )
+    return prompt
+
+
+def get_creative_seed_prompt() -> str:
+    prompt = (
+        "You are having a creative open-ended conversation inspired by the information in the context. "
+        "Please generate an open question for your conversation partner to start off the discussion. "
+        "Answer only with the question, do not include any other text."
+    )
+    return prompt
+
+
+def get_generic_seed_prompt() -> str:
+    return "Please generate a single chat message to begin a conversation about the information in the context. Make an open-ended request or provide a task."

src/ctx_to_lora/data/self_gen_template.py

@@ -0,0 +1,16 @@
+SELF_GEN_SYSTEM_MSG = "You are an honest and helpful assistant."
+
+SELF_QA_INTX = (
+    "# System Instruction\n"
+    "- The information provided is up-to-date information and/or the user instruction.\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"
+    "- If the provided information is an instruction, follow the instruction carefully.\n"
+    "\n---\n\n"
+    "# User Input\n"
+)
+
+PRE_CTX = "# Provided Information\n"
+
+QA_PROMPT_TEMPLATE = PRE_CTX + "{context}\n\n---\n\n" + SELF_QA_INTX + "{question}"
+PROMPT_TEMPLATE = "{context}\n\n{question}"

src/ctx_to_lora/eval_utils.py

@@ -4,6 +4,7 @@ import os
 import re
 import string
 import sys
+import time
 from argparse import Namespace
 from collections import Counter, defaultdict
 from dataclasses import fields
@@ -14,6 +15,7 @@ import pandas as pd
 import torch
 import yaml
 from datasets import disable_caching
+from peft import get_peft_model
 from transformers import (
     Seq2SeqTrainer,
     Seq2SeqTrainingArguments,
@@ -25,11 +27,13 @@ from transformers.trainer_utils import find_executable_batch_size
 from ctx_to_lora.data.collator import eval_collator, generation_collator
 from ctx_to_lora.data.definitions import (
     CLOSED_QA_DATASETS,
+    CTX_AFFIXES,
     LONGBENCH_E_TASKS,
     LONGBENCH_TASKS,
     MULTI_ANSWER_DATASETS,
 )
 from ctx_to_lora.data.processing import get_tokenized_dataset, load_answers
+from ctx_to_lora.data.self_gen_template import SELF_QA_INTX
 from ctx_to_lora.metrics import (
     LENGTH_BINS,
     Evaluator,
@@ -39,10 +43,18 @@ from ctx_to_lora.metrics import (
     compute_prefix_matching,
     compute_rouge,
 )
-from ctx_to_lora.model_loading import get_model, get_tokenizer
+from ctx_to_lora.model_loading import get_lora_config, get_model, get_tokenizer
 from ctx_to_lora.modeling import hypernet
+from ctx_to_lora.modeling.context_distillation import CtxDistillModel
 from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
-from ctx_to_lora.utils import clear_gpu, concat_list, setup_logging
+from ctx_to_lora.tracker.tracker import (
+    add_tracker,
+    print_global_tracker_stats,
+    print_tracker_stats,
+    reset_trackers,
+    save_tracker_stats_csv,
+)
+from ctx_to_lora.utils import clear_gpu, concat_list, get_run_name, setup_logging
 
 logger = logging.getLogger()
 
@@ -522,7 +534,13 @@ def decode_test_result(
         input_toks = sample["input_ids"][:start_idx]
 
         gen_toks = pred_toks[np.argmax(pred_toks != tokenizer.pad_token_id) :]
-        gen_toks = gen_toks[start_idx:]
+        # gen_toks = gen_toks[start_idx:]
+        suffix = np.array(CTX_AFFIXES[tokenizer.name_or_path]["suffix"])
+        # iterate over gen_toks and take the answer after the suffix
+        for i in range(len(gen_toks) - len(suffix), -1, -1):
+            if all(gen_toks[i : i + len(suffix)] == suffix):
+                gen_toks = gen_toks[i + len(suffix) :]
+                break
         gen_toks = np.where(gen_toks == -100, tokenizer.pad_token_id, gen_toks)
 
         d["input"] = tokenizer.decode(input_toks, skip_special_tokens=False)
@@ -543,7 +561,6 @@ def decode_test_result(
     return out
 
 
-@torch.inference_mode()
 def eval_generation(
     eval_trainer,
     tokenizer,
@@ -681,7 +698,6 @@ def eval_generation(
     return out
 
 
-@torch.no_grad()
 def eval_teacher_forcing(
     eval_trainer, datasets, split, remove_context
 ) -> dict[str, dict]:
@@ -736,6 +752,12 @@ def evaluate(
     ctx_name = None
     model_kwargs = dict(attn_implementation="flash_attention_2")
 
+    tokenizer = get_tokenizer(args.model_name_or_path, train=False)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    use_cd = False
+
     if model_name_or_path is None:
         try:
             state_dict = torch.load(checkpoint_path, weights_only=False)
@@ -749,6 +771,10 @@ def evaluate(
             use_flash_attn=True,
             use_sequence_packing=False,  # for generation
         )
+        add_tracker(model.base_model.generate, "generate")
+        add_tracker(model.generate_weights, "generate_weights")
+        add_tracker(model.combine_lora, "combine_lora")
+        add_tracker(model.apply_lora_to_layers, "apply_lora_to_layers")
     else:
         model = get_model(
             model_name_or_path,
@@ -757,12 +783,40 @@ def evaluate(
             model_kwargs=model_kwargs,
             use_flash_attn=True,
         )
+        add_tracker(model.generate, "generate")
+        if use_cd := getattr(args, "use_cd", False):
+            peft_config = get_lora_config(
+                model_name_or_path,
+                lora_r=8,
+                lora_dropout=0,
+                target_modules=["down_proj"],
+            )
+            peft_config.lora_alpha = 16
+            peft_model = get_peft_model(model, peft_config)
+            q_model = get_model(
+                "google/gemma-3-4b-it", train=False, requires_grad=False
+            )
+            model = CtxDistillModel(
+                peft_model,
+                prefix_tokens=torch.tensor(
+                    CTX_AFFIXES[model_name_or_path]["prefix"], device=model.device
+                ),
+                pad_token_id=tokenizer.pad_token_id,
+                update_iterations=args.cd_update_iterations,
+                q_model=q_model,  # peft_model if args.cd_use_gen_q else None,
+                num_gen_q=args.num_gen_q,
+                tokenizer=tokenizer,
+            )
+            add_tracker(model._distill_context, "distill_context")
+            add_tracker(model.generate_questions, "generate_questions")
+            add_tracker(model.teacher_generate, "teacher_generate")
+            add_tracker(model.student_generate, "student_generate")
 
-    tokenizer = get_tokenizer(args.model_name_or_path, train=False)
-    if tokenizer.pad_token_id is None:
-        tokenizer.pad_token_id = tokenizer.eos_token_id
     base_model = (
-        model.base_model if isinstance(model, ModulatedPretrainedModel) else model
+        model.base_model
+        if isinstance(model, ModulatedPretrainedModel)
+        or isinstance(model, CtxDistillModel)
+        else model
     )
     base_model.config.pad_token_id = tokenizer.pad_token_id
     base_model.generation_config.pad_token_id = tokenizer.pad_token_id
@@ -775,7 +829,7 @@ def evaluate(
 
     add_ctx_to_chat = (
         not isinstance(model, ModulatedPretrainedModel) and not args.remove_context
-    )
+    ) or isinstance(model, CtxDistillModel)
     ctx_model_max_len = (
         model.ctx_encoder.config.max_position_embeddings
         if isinstance(model, ModulatedPretrainedModel)
@@ -798,6 +852,7 @@ def evaluate(
         use_kl_loss=False,
         max_new_tokens=max_new_tokens,
         set_format="pt",
+        add_self_distill_template=use_cd,  # only for eval
     )
 
     datasets = dict()
@@ -868,6 +923,20 @@ def evaluate(
     if max_ctx_chunk_len > 0:
         model.generate = model.generate_with_multi_loras
 
+    if isinstance(model, CtxDistillModel):
+        sep_seq = (
+            tokenizer(
+                SELF_QA_INTX.strip("\n"), add_special_tokens=False, return_tensors="pt"
+            )
+            .input_ids[0]
+            .to(model.device)
+        )
+        model.generate = partial(
+            model.generate,
+            ctx_inp_sep_seq=sep_seq,
+            reset=True,
+        )
+
     trainer_kwargs = {
         "model": model,
         "args": eval_trainer_args,
@@ -895,17 +964,26 @@ def evaluate(
             out.update(metrics)
     else:
         eval_trainer = CustomSeq2SeqTrainer(**trainer_kwargs)
-        metrics = eval_generation(
-            eval_trainer,
-            tokenizer,
-            ctx_tokenizer,
-            datasets,
-            answers,
-            split,
-            args.remove_context,
-            gen_kwargs,
-        )
-        out.update(metrics)
+        for ds_name, ds in datasets.items():
+            metrics = eval_generation(
+                eval_trainer,
+                tokenizer,
+                ctx_tokenizer,
+                {ds_name: ds},
+                answers,
+                split,
+                args.remove_context,
+                gen_kwargs,
+            )
+            out.update(metrics)
+
+            print_tracker_stats()
+            print_global_tracker_stats()
+            ds_suffix = "_no_context" if args.remove_context else ""
+            save_tracker_stats_csv(
+                f"{args.logging_dir}/{split}_{ds_name}{ds_suffix}_tracked_stats.csv"
+            )
+            reset_trackers()
 
     clear_gpu()
     return out
@@ -922,11 +1000,17 @@ def run_eval(
     remove_context: bool = False,
     max_new_tokens: int = 256,
     generative: bool = False,
+    use_cd: bool = False,
+    cd_update_iterations: int = 10,
+    cd_use_gen_q: bool = False,
+    num_gen_q: int = 20,
 ) -> None:
     """Run evaluation with the specified parameters."""
     assert bool(model_name_or_path) ^ bool(checkpoint_path), (
         "Either --model_name_or_path or --checkpoint_path must be provided"
     )
+    if use_cd and eval_batch_size != 1:
+        raise ValueError("When using context distillation, eval_batch_size must be 1.")
 
     disable_caching()
     set_seed(42)
@@ -943,6 +1027,9 @@ def run_eval(
     torch.backends.cuda.matmul.allow_tf32 = False
     torch.backends.cudnn.allow_tf32 = False
 
+    slurm_job_id = f"_{os.getenv('SLURM_JOB_ID')}" if os.getenv("SLURM_JOB_ID") else ""
+    run_name = get_run_name(seed_str=time.strftime("%Y%m%d-%H%M%S") + slurm_job_id)
+
     if checkpoint_path:
         checkpoint_dir = "/".join(checkpoint_path.split("/")[:-1])
         run_dir = "/".join(checkpoint_path.split("/")[:-2])
@@ -954,8 +1041,8 @@ def run_eval(
         print(f"checkpoint_path: {checkpoint_path}")
         print(f"run_dir: {run_dir}")
 
-        args.output_dir = f"{run_dir}/eval-results-{cur_it}"
-        args.logging_dir = f"{run_dir}/eval-results-{cur_it}"
+        args.output_dir = f"{run_dir}/eval-results-{cur_it}/{run_name}"
+        args.logging_dir = f"{run_dir}/eval-results-{cur_it}/{run_name}"
         args.run_name = run_dir.split("/")[-1]
         # modulated model doesn't see ctx by default
         # but remove_context has to be false for correct file naming
@@ -963,16 +1050,22 @@ def run_eval(
     else:
         args = Namespace(
             model_name_or_path=model_name_or_path,
-            output_dir=f"eval_results/{model_name_or_path}",
-            logging_dir=f"eval_results/{model_name_or_path}",
-            run_name=f"eval_results/{model_name_or_path}",
+            output_dir=f"eval_results/{model_name_or_path}/{run_name}",
+            logging_dir=f"eval_results/{model_name_or_path}/{run_name}",
+            run_name=f"eval_results/{model_name_or_path}/{run_name}",
             val_ds_names=[],
             test_ds_names=[],
             remove_context=remove_context,
         )
+        if use_cd:
+            args.use_cd = use_cd
+            args.cd_update_iterations = cd_update_iterations
+            args.cd_use_gen_q = cd_use_gen_q
+            args.num_gen_q = num_gen_q
     if max_val_samples_per_ds > 0:
         args.max_val_samples_per_ds = max_val_samples_per_ds
     setup_logging(args.logging_dir)
+    logger.debug(f"CMD: {' '.join(os.sys.argv)}")
 
     # Override dataset names if provided via CLI
     if datasets:

src/ctx_to_lora/modeling/context_distillation.py

@@ -0,0 +1,503 @@
+import logging
+import random
+from typing import Any
+
+import torch
+from jaxtyping import Integer
+from peft import PeftModel
+from peft.tuners.tuners_utils import BaseTunerLayer, check_target_module_exists
+from torch import Tensor, nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+from ctx_to_lora.data.definitions import CTX_AFFIXES
+from ctx_to_lora.data.q_generation_prompts import (
+    PROMPT_TEMPLATE,
+    SYSTEM_PROMPT,
+    get_creative_seed_prompt,
+    get_generic_seed_prompt,
+    get_question_seed_prompt,
+    get_use_case_seed_prompt,
+)
+from ctx_to_lora.data.self_gen_template import SELF_QA_INTX
+from ctx_to_lora.utils import log_num_train_params
+
+logger = logging.getLogger()
+
+
+def build_messages(
+    context: str,
+    question_weight: float,
+    use_case_weight: float,
+    creative_weight: float,
+    generic_weight: float,
+) -> list[dict]:
+    # Create list of instruction functions with their weights
+    instruction_options = [
+        ("question", get_question_seed_prompt, question_weight),
+        ("use_case", get_use_case_seed_prompt, use_case_weight),
+        ("creative", get_creative_seed_prompt, creative_weight),
+        ("generic", get_generic_seed_prompt, generic_weight),
+    ]
+
+    q_types, functions, weights = zip(*instruction_options)
+
+    chosen_idx = random.choices(range(len(instruction_options)), weights=weights)[0]
+
+    function = functions[chosen_idx]
+
+    custom_instruction = function()
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {
+            "role": "user",
+            "content": PROMPT_TEMPLATE.format(
+                context=context, instruction=custom_instruction
+            ),
+        },
+    ]
+    return messages
+
+
+def get_shifted_label_pos(labels):
+    pos = torch.where(labels != -100)
+    # (batch_idx, token_idx)
+    return (pos[0], pos[1] - 1)
+
+
+def logits_at_positions(outputs: ModelOutput, pos) -> Tensor:
+    logits = outputs.logits
+    return logits[pos[0], pos[1]]
+
+
+def ctx_inp_split(
+    ctx_inp_ids, ctx_inp_sep_seq, pad_token_id, prefix_tokens=None, padding_side="right"
+):
+    # Split each row in ctx_inp_ids at the first occurrence of ctx_inp_sep_seq
+    # Return the part after the separator for each row
+    batch_size = ctx_inp_ids.size(0)
+    sep_len = ctx_inp_sep_seq.size(0)
+    out_inp = []
+    out_ctx = []
+    for i in range(batch_size):
+        row = ctx_inp_ids[i]
+        # Find where the separator starts
+        for j in range(row.size(0) - sep_len + 1):
+            if torch.equal(row[j : j + sep_len], ctx_inp_sep_seq):
+                out_ctx.append(row[:j])
+                if prefix_tokens is not None:
+                    out_inp.append(
+                        torch.cat([prefix_tokens, row[j + sep_len :]], axis=-1)
+                    )
+                else:
+                    out_inp.append(row[j + sep_len :])
+                break
+        else:
+            # If separator not found
+            raise ValueError(f"Separator sequence not found in row {i}")
+    out_inp = torch.nn.utils.rnn.pad_sequence(
+        out_inp, batch_first=True, padding_value=pad_token_id, padding_side=padding_side
+    )
+    out_ctx = torch.nn.utils.rnn.pad_sequence(
+        out_ctx, batch_first=True, padding_value=pad_token_id, padding_side=padding_side
+    )
+    return out_ctx, out_inp
+
+
+def get_peft_layers(model, peft_config):
+    out = []
+    for module_name, module in model.named_modules():
+        if not check_target_module_exists(peft_config, module_name):
+            continue
+        if not isinstance(module, BaseTunerLayer):
+            continue
+        # support just Linear layer for now
+        # all modules should be a leave module that is Linear layer
+        assert isinstance(module.base_layer, nn.Linear), (
+            "all modules should be a leave module that is Linear layer"
+        )
+
+        # this should always pass
+        name = module_name.split(".")[-1]
+        assert name in peft_config.target_modules
+        out.append(module)
+
+    return out
+
+
+class CtxDistillModel(nn.Module):
+    def __init__(
+        self,
+        base_model: PeftModel,
+        prefix_tokens: Integer[Tensor, "n"],
+        pad_token_id: int,
+        update_iterations: int,
+        q_model: PreTrainedModel | None = None,
+        num_gen_q: int | None = None,
+        tokenizer=None,
+    ):
+        super().__init__()
+        self.register_module("base_model", base_model)
+        self.register_module("q_model", q_model)
+        if q_model is not None:
+            assert num_gen_q is not None, (
+                "num_gen_q must be provided if q_model is used"
+            )
+            self.num_gen_q = num_gen_q
+        self.register_buffer("prefix_tokens", prefix_tokens)
+        self.tokenizer = tokenizer
+        self.pad_token_id = pad_token_id
+        self.update_iterations = update_iterations
+        self.device = base_model.device
+        self.to(self.device)
+
+        self.peft_config = base_model.peft_config["default"]
+        self.adapter_name = "default"
+        self.base_model.set_adapter("default")
+        for layer in get_peft_layers(self.base_model, self.peft_config):
+            for name, p in layer.named_parameters():
+                if "lora_A" in name or "lora_B" in name:
+                    p.requires_grad = True
+        log_num_train_params(self.base_model)
+        self._init_optim()
+
+    @property
+    def generation_config(self):
+        return self.base_model.generation_config
+
+    def _init_optim(self):
+        self.optimizer = torch.optim.AdamW(
+            [
+                p
+                for l in get_peft_layers(self.base_model, self.peft_config)
+                for p in l.parameters()
+                if p.requires_grad
+            ],
+            lr=1e-4,
+        )
+
+    def reset_lora(self):
+        print("Resetiing LoRA")
+        for layer in get_peft_layers(self.base_model, self.peft_config):
+            layer.reset_lora_parameters(self.adapter_name, init_lora_weights=True)
+        self._init_optim()
+
+    @torch.enable_grad()
+    def _distill_context(
+        self,
+        ctx_inp_res_ids: Integer[Tensor, "bs ctx_inp_length"],
+        ctx_inp_res_attention_mask: Integer[Tensor, "bs ctx_inp_length"],
+        teacher_labels: Integer[Tensor, "bs ctx_inp_length"],
+        inp_res_ids: Integer[Tensor, "bs inp_length"],
+        inp_res_attention_mask: Integer[Tensor, "bs inp_length"],
+        student_labels: Integer[Tensor, "bs inp_length"],
+    ):
+        # Implements KD-style loss by computing teacher (with context) and student (no context) log-probs locally.
+
+        teacher_shifted_label_pos = get_shifted_label_pos(teacher_labels)
+        student_shifted_label_pos = get_shifted_label_pos(student_labels)
+
+        with torch.no_grad():
+            teacher_outputs = self.base_model(
+                ctx_inp_res_ids, attention_mask=ctx_inp_res_attention_mask
+            )
+            teacher_logits = logits_at_positions(
+                teacher_outputs, teacher_shifted_label_pos
+            )
+            # Use top-k teacher probabilities only
+            K = 16
+            topk_vals, topk_idx = teacher_logits.topk(K, dim=-1)
+            teacher_denom = torch.logsumexp(
+                teacher_logits.float(), dim=-1, keepdim=True
+            )
+            teacher_p = (topk_vals - teacher_denom).exp().detach()  # [N, K]
+
+        # Optimization loop: match student to teacher using selected indices
+        was_training = self.training
+        self.train()
+
+        print(f"Starting context distillation for {self.update_iterations} iterations")
+
+        for iteration in range(self.update_iterations):
+            self.optimizer.zero_grad()
+            student_outputs = self.base_model(
+                inp_res_ids, attention_mask=inp_res_attention_mask
+            )
+            student_logits = logits_at_positions(
+                student_outputs, student_shifted_label_pos
+            )
+            student_denom = torch.logsumexp(
+                student_logits.float(), dim=-1, keepdim=True
+            )
+            selected_student_logits = student_logits.gather(-1, topk_idx)
+            student_logq = selected_student_logits - student_denom  # [N, K]
+            token_losses = -(teacher_p * student_logq).sum(dim=-1)  # [N]
+            loss = token_losses.mean()
+
+            # Log metrics for this iteration
+            print(
+                f"Iteration {iteration + 1}/{self.update_iterations}: "
+                f"Loss={loss.item():.4f}"
+            )
+
+            loss.backward()
+            self.optimizer.step()
+
+        if not was_training:
+            self.eval()
+
+    def generate_questions(self, *args, **kwargs):
+        return self.q_model.generate(*args, **kwargs)
+
+    def teacher_generate(self, *args, **kwargs):
+        # rename for separate timing
+        return self.base_model.generate(*args, **kwargs)
+
+    def student_generate(self, *args, **kwargs):
+        # rename for separate timing
+        return self.base_model.generate(*args, **kwargs)
+
+    def generate(
+        self,
+        # ctx_ids: Integer[Tensor, "n_chunks ctx_length"] | None = None,
+        # ctx_attn_mask: Integer[Tensor, "n_chunks ctx_length"] | None = None,
+        # ctx_position_ids: Integer[Tensor, "n_chunks ctx_length"] | None = None,
+        # n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None,
+        # n_queries: Integer[Tensor, "n_ctx"] | None = None,
+        *model_inputs_args: Any,
+        ctx_inp_sep_seq: Integer[Tensor, "l"],
+        reset: bool,
+        **model_inputs_kwargs: dict[str, Any],
+    ):
+        # where to get the questions???
+        # TODO: contex indep variant
+        # self.generate_questions(ctx_ids)
+
+        # update peft module with CD (if labels provided)
+
+        if reset:
+            self.reset_lora()
+
+        # teacher tokens
+        if model_inputs_args:
+            ctx_inp_ids = model_inputs_args[0]
+        else:
+            ctx_inp_ids = model_inputs_kwargs.pop("input_ids")
+        _, orig_inp_ids = ctx_inp_split(
+            ctx_inp_ids,
+            ctx_inp_sep_seq,
+            self.pad_token_id,
+            self.prefix_tokens,
+            padding_side="left",
+        )
+        ctx_inp_attention_mask = model_inputs_kwargs.pop("attention_mask")
+
+        if self.q_model is not None:
+            # Extract context-only portion after separator (remove prefix tokens from first row)
+            ctx_ids_full, _ = ctx_inp_split(
+                ctx_inp_ids, ctx_inp_sep_seq, self.pad_token_id
+            )  # [bs, var_len]
+            ctx_ids = ctx_ids_full[0, len(self.prefix_tokens) :]
+            ctx_txt = self.tokenizer.decode(ctx_ids, skip_special_tokens=True)
+            # Build multiple instruction variants
+            messages_list = [
+                build_messages(ctx_txt, 1, 1, 1, 1) for _ in range(self.num_gen_q)
+            ]
+            q_inputs = self.tokenizer.apply_chat_template(
+                messages_list,
+                tokenize=True,
+                add_special_tokens=False,
+                padding=True,
+                truncation=False,
+                add_generation_prompt=True,
+                return_dict=True,
+                return_tensors="pt",
+            )
+            q_inputs = {k: v.to(self.device) for k, v in q_inputs.items()}
+
+            with torch.no_grad():
+                question_outputs = self.generate_questions(
+                    input_ids=q_inputs["input_ids"],
+                    attention_mask=q_inputs.get("attention_mask", None),
+                    max_new_tokens=256,
+                    do_sample=True,
+                    top_p=0.95,
+                    temperature=1.0,  # high temp for diverse questions
+                )
+            # Slice off the prompt portion
+            gen_only = question_outputs[:, q_inputs["input_ids"].shape[-1] :]
+            questions = self.tokenizer.batch_decode(gen_only, skip_special_tokens=True)
+            questions = [q.split("Message:")[-1].strip() for q in questions]
+
+            ctx_inp_messages = [
+                [{"role": "user", "content": f"{ctx_txt}\n\n{SELF_QA_INTX}\n\n{q}"}]
+                for q in questions
+            ]
+            encoded_ctx_inp = self.tokenizer.apply_chat_template(
+                ctx_inp_messages,
+                tokenize=True,
+                add_special_tokens=False,
+                padding=True,
+                truncation=False,
+                add_generation_prompt=True,
+                return_dict=True,
+                return_tensors="pt",
+            )
+            encoded_ctx_inp = {k: v.to(self.device) for k, v in encoded_ctx_inp.items()}
+            ctx_inp_ids = encoded_ctx_inp["input_ids"]
+            ctx_inp_attention_mask = encoded_ctx_inp["attention_mask"]
+
+        # TODO: check labels + loss calculation + padding
+
+        # sample responses first
+        ctx_inp_res_ids = self.teacher_generate(
+            ctx_inp_ids,
+            attention_mask=ctx_inp_attention_mask,
+            **model_inputs_kwargs,
+        )
+        ctx_inp_res_attention_mask = torch.where(
+            ctx_inp_res_ids != self.pad_token_id, 1, 0
+        ).long()
+        ctx_inp_res_txt = self.tokenizer.batch_decode(ctx_inp_res_ids)
+
+        bs = ctx_inp_ids.shape[0]
+        res_len = ctx_inp_res_ids.shape[-1] - ctx_inp_ids.shape[-1]
+        res_ids = ctx_inp_res_ids[:, -res_len:]  # correct
+
+        pads = torch.full_like(ctx_inp_ids, self.pad_token_id)
+        teacher_labels = torch.cat([pads, res_ids], dim=-1)
+        teacher_labels = torch.where(
+            teacher_labels != self.pad_token_id, teacher_labels, -100
+        )
+
+        # student tokens
+        _, inp_res_ids = ctx_inp_split(
+            ctx_inp_res_ids,
+            ctx_inp_sep_seq,
+            self.pad_token_id,
+            self.prefix_tokens,
+            padding_side="left",
+        )
+        inp_res_attention_mask = torch.where(
+            inp_res_ids != self.pad_token_id, 1, 0
+        ).long()
+
+        # _, inp_ids = ctx_inp_split(
+        #     ctx_inp_ids,
+        #     ctx_inp_sep_seq.to(ctx_inp_res_ids.device),
+        #     self.pad_token_id,
+        #     self.prefix_tokens,
+        # )
+
+        student_labels = inp_res_ids.clone()
+        student_labels[:, :-res_len] = -100
+        student_labels = torch.where(
+            student_labels != self.pad_token_id, student_labels, -100
+        )
+
+        self._distill_context(
+            ctx_inp_res_ids,
+            ctx_inp_res_attention_mask,
+            teacher_labels,
+            inp_res_ids,
+            inp_res_attention_mask,
+            student_labels,
+        )
+
+        # _, inp_ids = ctx_inp_split(
+        #     ctx_inp_ids,
+        #     ctx_inp_sep_seq.to(ctx_inp_res_ids.device),
+        #     self.pad_token_id,
+        #     self.prefix_tokens,
+        #     padding_side="left",
+        # )
+
+        # # inp_ids = torch.cat([self.prefix_tokens.expand(bs, -1), inp_ids], dim=-1)
+        # # inp_ids = inp_res_ids[:, :-res_len]
+        # print(self.tokenizer.batch_decode(inp_ids))
+        inp_attention_mask = torch.where(orig_inp_ids != self.pad_token_id, 1, 0).long()
+        model_inputs_kwargs.pop("attention_mask", None)
+        model_inputs_kwargs.pop("input_ids", None)
+        model_outputs = self.student_generate(
+            orig_inp_ids, attention_mask=inp_attention_mask, **model_inputs_kwargs
+        )
+        return model_outputs
+
+
+if __name__ == "__main__":
+    from ctx_to_lora.data.processing import load_and_process_dataset
+    from ctx_to_lora.model_loading import get_lora_config, get_model_and_tokenizer
+
+    model_name = "google/gemma-2-2b-it"
+    peft_config = get_lora_config(
+        model_name, r=8, target_modules=["down_proj"], lora_dropout=0.0
+    )
+    peft_config.lora_alpha = 16
+    model, tokenizer = get_model_and_tokenizer(
+        model_name,
+        train=False,
+        requires_grad=False,
+        peft_config=peft_config,
+    )
+
+    ds = load_and_process_dataset("pwc", split="train", num_proc=8)
+    ctx = ds[0]["context"]
+    inp = ds[1]["prompts"][0]
+    # Build a simple context/input pair separated by a unique token sequence
+    sep_text = SELF_QA_INTX
+    # ctx = "# Provided Information\nMy name is Tan."
+    # ctx = "# Provided Info"
+    # inp = "What is my name? Reply with my name only. No explanation."
+    prompt = f"{ctx}\n\n{sep_text}\n\n{inp}"
+    messages = [{"role": "user", "content": prompt}]
+    encoded = tokenizer.apply_chat_template(
+        messages, return_tensors="pt", return_dict=True
+    )
+    # encoded = tokenizer(prompt, return_tensors="pt")
+    encoded = {k: v.to(model.device) for k, v in encoded.items()}
+
+    prefix_tokens = CTX_AFFIXES[model_name]["prefix"]
+    prefix_tokens = torch.tensor(prefix_tokens, dtype=torch.long)
+
+    cd_model = CtxDistillModel(
+        base_model=model,
+        prefix_tokens=prefix_tokens,
+        pad_token_id=tokenizer.pad_token_id,
+        update_iterations=200,
+        q_model=model,
+        num_gen_q=20,
+        tokenizer=tokenizer,
+    )
+
+    sep_ids = (
+        tokenizer(sep_text.strip("\n"), add_special_tokens=False, return_tensors="pt")
+        .input_ids[0]
+        .to(model.device)
+    )
+
+    with torch.no_grad():
+        for _ in range(1):
+            base_model_res = model.generate(
+                encoded["input_ids"],
+                attention_mask=encoded["attention_mask"],
+                max_new_tokens=256,
+                do_sample=False,
+            )
+            print(
+                f"Base model response:{tokenizer.batch_decode(base_model_res, skip_special_tokens=False)}"
+            )
+
+            outputs = cd_model.generate(
+                encoded["input_ids"],
+                attention_mask=encoded["attention_mask"],
+                ctx_inp_sep_seq=sep_ids,
+                reset=True,
+                max_new_tokens=256,
+                do_sample=False,
+            )
+            # print(
+            #     f"Teacher response:{tokenizer.batch_decode(teacher_outputs, skip_special_tokens=False)}"
+            # )
+            print(
+                f"Student response: {tokenizer.batch_decode(outputs, skip_special_tokens=False)}"
+            )

src/ctx_to_lora/modeling/hypernet.py

@@ -1103,6 +1103,14 @@ class ModulatedPretrainedModel(nn.Module):
         )
         return model_outputs
 
+    def combine_lora(self, *args, **kwargs):
+        # for timing
+        return combine_lora(*args, **kwargs)
+
+    def apply_lora_to_layers(self, *args, **kwargs):
+        # for timing
+        return apply_lora_to_layers(*args, **kwargs)
+
     @torch.inference_mode()
     def generate(
         self,
@@ -1159,7 +1167,7 @@ class ModulatedPretrainedModel(nn.Module):
             )
 
         if generated_loras is not None:
-            generated_loras = combine_lora(
+            generated_loras = self.combine_lora(
                 generated_loras,
                 n_ctx_chunks,
                 lora_bias=self.hypernet.get_head_bias(),

src/ctx_to_lora/modeling/test_ctx.txt

@@ -0,0 +1,15 @@
+Actress Halle Berry has been sharing a number of stunning photos from the time she has spent in Morocco and she just posted a new one to her Instagram page that fans will not want to miss. The Sahara Desert is a stunning backdrop for any photo shoot, but when Berry is involved, there is no question that the beauty of the area is elevated significantly.
+
+In Halle Berry's newest Instagram post, she shared a photo showing her from the back as she looks out at the Sahara Desert. It's a rather sultry shot, as Halle's back is entirely bare. Berry is sitting on the sand barefoot, and she has one arm holding up the black fabric she has around her.
+
+Berry's other arm is sitting on her knees and she is leaning her head against her hand. Halle's long, dark tresses are free and wavy, cascading down her back, and her face isn't visible at all from this angle. Other than the black fabric draping Halle's fit physique, the only other thing she appears to be wearing is a turquoise ring or something similar.
+
+The actress kept her caption simple, yet intriguing. Berry said that everything has changed, but at the same time, she's "more me than I've ever been." While Halle's fans would probably love more insight or context into the phrasing she used, it somehow synced perfectly with the sultry photo she shared.
+
+Halle has about 4 million followers on Instagram, and this latest post quickly garnered about 30,000 likes and several hundred comments. The picture definitely resonated with people and many noted that she was absolutely beautiful in the stunning shot.
+
+Other recent tidbits from Berry's Instagram teased a reference to her character of Sofia from the upcoming movie John Wick: Chapter 3, as well as a tribute to comic great Stan Lee after his passing earlier this week. Both Halle and her John Wick co-star Keanu Reeves have been filming in Morocco for a bit now and it looks like the setting has been inspirational for the actress.
+
+Berry is now 52-years-old, but as the Inquisitr has previously shared, her trainer says she has the body of a 25-year-old. Halle makes staying fit a top priority for her and it certainly shows in this sultry shot she just shared. The actress has endured some difficult times in her personal life over the past few years, but it looks like both physically and mentally she's in the best shape of her life.
+
+Halle Berry's time in Morocco has been inspirational not only to her personally but to many of her followers and they cannot wait to see what she shares with them next.

src/ctx_to_lora/tracker/cuda_memory_tracker.py

@@ -0,0 +1,353 @@
+"""Lightweight CUDA memory tracking utilities for measuring per-method peak memory usage.
+
+Usage:
+    x = SomeClass()
+    add_memory_tracker(x.some_method, "some_method")  # wraps the bound method in-place
+    x.some_method(...)
+    print_aggregate_memory_stats("some_method")
+
+Design notes:
+- Mirrors the API of timer.py but records CUDA memory (peak increase in bytes) per call.
+- add_memory_tracker mutates the instance method with a wrapper (idempotent: double wrap avoided).
+- Global registry: { name: [int, ...] } storing per-call peak memory increase (bytes).
+- print_aggregate_memory_stats prints summary stats (count, total, mean, median, min, max, p95, last).
+- If CUDA or torch is unavailable, wrappers degrade gracefully (no measurements recorded).
+
+Metrics collected per call (if CUDA available):
+- peak_increase_bytes: (torch.cuda.max_memory_allocated() - start_allocated)
+  This captures the maximum additional memory pressure during the call.
+
+Caveats:
+- Rapid allocate/free patterns entirely inside the call still reflect peak transient usage.
+- Asynchronous CUDA ops: we synchronize before and after to improve accuracy. This may
+  slightly affect performance timings but is necessary for memory correctness.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from statistics import mean, median, stdev
+from typing import Any
+
+try:  # Optional dependency handling
+    import torch  # type: ignore
+except Exception:  # pragma: no cover - torch absence path
+    torch = None  # type: ignore
+
+# Global memory registry: name -> list of peak memory increases (bytes)
+MEMORY_REGISTRY: dict[str, list[int]] = {}
+
+
+def _cuda_available() -> bool:
+    return bool(torch is not None and torch.cuda.is_available())
+
+
+def add_memory_tracker(func: Callable, name: str) -> None:
+    """Attach a CUDA memory tracking wrapper to a bound method.
+
+    Parameters
+    ----------
+    func : Callable
+        A *bound* instance method (instance.method). Raises ValueError if unbound.
+    name : str
+        Key under which memory stats are recorded in MEMORY_REGISTRY.
+    """
+    if not hasattr(func, "__self__") or getattr(func, "__self__") is None:
+        if getattr(func, "__is_memory_wrapper__", False):  # already wrapped
+            return
+        raise ValueError(
+            "add_memory_tracker expects a bound method: call with instance.method"
+        )
+
+    instance = func.__self__
+    method_name = getattr(func, "__name__", None)
+    if method_name is None:
+        raise ValueError("Cannot determine method name for provided callable")
+
+    existing = getattr(instance, method_name, None)
+    if getattr(existing, "__is_memory_wrapper__", False):  # idempotent
+        return
+
+    orig_bound = func
+
+    def tracked(*args: Any, **kwargs: Any):  # noqa: D401 - simple wrapper
+        if not _cuda_available():
+            return orig_bound(*args, **kwargs)
+        # Synchronize to get a clean baseline
+        torch.cuda.synchronize()
+        torch.cuda.reset_peak_memory_stats()
+        start_alloc = torch.cuda.memory_allocated()
+        try:
+            return orig_bound(*args, **kwargs)
+        finally:
+            torch.cuda.synchronize()
+            peak_alloc = torch.cuda.max_memory_allocated()
+            peak_increase = peak_alloc - start_alloc
+            # Record only if positive (avoid negative due to potential race, though improbable)
+            if peak_increase < 0:
+                peak_increase = 0
+            MEMORY_REGISTRY.setdefault(name, []).append(int(peak_increase))
+
+    tracked.__name__ = method_name
+    tracked.__doc__ = getattr(orig_bound, "__doc__")
+    tracked.__qualname__ = getattr(orig_bound, "__qualname__", method_name)
+    tracked.__is_memory_wrapper__ = True  # type: ignore[attr-defined]
+    tracked.__wrapped__ = orig_bound  # type: ignore[attr-defined]
+    tracked.__memory_name__ = name  # type: ignore[attr-defined]
+
+    setattr(instance, method_name, tracked)
+
+
+def _format_bytes(num_bytes: float) -> str:
+    """Human-readable byte formatting (base-2)."""
+    if num_bytes < 1024:
+        return f"{int(num_bytes):5d}B"
+    units = ["KiB", "MiB", "GiB", "TiB"]
+    value = float(num_bytes)
+    for u in units:
+        value /= 1024.0
+        if value < 1024.0:
+            return f"{value:8.3f}{u}"
+    return f"{value:8.3f}PiB"  # Extremely unlikely
+
+
+def compute_aggregate_memory_stats(
+    name: str | None = None,
+) -> dict[str, dict[str, float]] | None:
+    """Compute aggregate CUDA memory statistics for specific trackers.
+
+    Parameters
+    ----------
+    name : Optional[str]
+        Specific tracker name to compute stats for. If None, all trackers are computed.
+
+    Returns
+    -------
+    Optional[Dict[str, Dict[str, float]]]
+        None if no data, else dict mapping tracker names to their stats.
+        Each stats dict contains: count, total, mean, median, min, max, p95, last, std.
+    """
+    if not MEMORY_REGISTRY:
+        return None
+
+    keys = [name] if name else sorted(MEMORY_REGISTRY.keys())
+    valid_keys = [k for k in keys if k in MEMORY_REGISTRY and MEMORY_REGISTRY[k]]
+    if not valid_keys:
+        return None
+
+    result = {}
+    for k in valid_keys:
+        data = MEMORY_REGISTRY[k]
+        data_sorted = sorted(data)
+        cnt = len(data)
+        total = sum(data)
+        avg = mean(data)
+        med = median(data)
+        std = stdev(data) if cnt > 1 else 0.0
+        mn = data_sorted[0]
+        mx = data_sorted[-1]
+        p95_index = int(0.95 * (cnt - 1))
+        p95 = data_sorted[p95_index]
+        last = data[-1]
+        result[k] = {
+            "count": float(cnt),
+            "total": float(total),
+            "mean": float(avg),
+            "median": float(med),
+            "min": float(mn),
+            "max": float(mx),
+            "p95": float(p95),
+            "last": float(last),
+            "std": float(std),
+        }
+    return result
+
+
+def save_memory_stats_csv(file_path: str, name: str | None = None) -> None:
+    """Save aggregate CUDA memory statistics to a CSV file.
+
+    Parameters
+    ----------
+    file_path : str
+        Path where the CSV file will be saved.
+    name : Optional[str]
+        Specific tracker name to export. If None, all trackers are exported.
+    """
+    import csv
+
+    stats = compute_aggregate_memory_stats(name)
+    if stats is None:
+        raise ValueError("No memory data available to export")
+
+    with open(file_path, "w", newline="") as csvfile:
+        fieldnames = [
+            "name",
+            "count",
+            "total",
+            "mean",
+            "median",
+            "min",
+            "max",
+            "p95",
+            "last",
+            "std",
+        ]
+        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+        writer.writeheader()
+
+        for tracker_name, data in stats.items():
+            row = {"name": tracker_name}
+            row.update(data)
+            writer.writerow(row)
+
+
+def print_aggregate_memory_stats(name: str | None = None) -> None:
+    """Print aggregate CUDA memory stats for one or all tracked names.
+
+    Parameters
+    ----------
+    name : Optional[str]
+        Specific name to report; if None, report all.
+    """
+    stats = compute_aggregate_memory_stats(name)
+    if stats is None:
+        print("[mem] No memory data collected.")
+        return
+
+    keys = [name] if name else sorted(MEMORY_REGISTRY.keys())
+    missing = [k for k in keys if k not in MEMORY_REGISTRY or not MEMORY_REGISTRY[k]]
+    if missing:
+        print(f"[mem] No data for: {', '.join(missing)}")
+
+    if not stats:
+        return
+
+    header = (
+        f"{'name':20} {'count':>6} {'total':>12} {'mean':>12} {'median':>12} "
+        f"{'min':>12} {'max':>12} {'p95':>12} {'std':>12} {'last':>12}"
+    )
+    print(header)
+    print("-" * len(header))
+
+    for k, data in stats.items():
+        print(
+            f"{k:20} {int(data['count']):6d} {_format_bytes(data['total']):>12} "
+            f"{_format_bytes(data['mean']):>12} {_format_bytes(data['median']):>12} "
+            f"{_format_bytes(data['min']):>12} {_format_bytes(data['max']):>12} "
+            f"{_format_bytes(data['p95']):>12} {_format_bytes(data['std']):>12} "
+            f"{_format_bytes(data['last']):>12}"
+        )
+
+
+def compute_global_memory_stats() -> dict[str, float] | None:
+    """Compute aggregate stats across all recorded memory entries.
+
+    Returns
+    -------
+    Optional[Dict[str, float]]
+        None if no data; else dict with count, total, mean, median, min, max, p95, std.
+    """
+    if not MEMORY_REGISTRY:
+        return None
+    all_values: list[int] = []
+    for lst in MEMORY_REGISTRY.values():
+        all_values.extend(lst)
+    if not all_values:
+        return None
+    data_sorted = sorted(all_values)
+    cnt = len(all_values)
+    total = float(sum(all_values))
+    avg = mean(all_values)
+    med = median(all_values)
+    std = stdev(all_values) if cnt > 1 else 0.0
+    mn = float(data_sorted[0])
+    mx = float(data_sorted[-1])
+    p95_index = int(0.95 * (cnt - 1))
+    p95 = float(data_sorted[p95_index])
+    return {
+        "count": float(cnt),
+        "total": total,
+        "mean": float(avg),
+        "median": float(med),
+        "min": mn,
+        "max": mx,
+        "p95": p95,
+        "std": float(std),
+    }
+
+
+def print_global_memory_stats() -> None:
+    """Pretty-print global stats across all memory registry entries."""
+    stats = compute_global_memory_stats()
+    if stats is None:
+        print("[mem] No memory data collected.")
+        return
+    header = (
+        f"{'scope':20} {'count':>6} {'total':>12} {'mean':>12} {'median':>12} "
+        f"{'min':>12} {'max':>12} {'p95':>12} {'std':>12}"
+    )
+    print(header)
+    print("-" * len(header))
+    print(
+        f"{'<ALL>':20} {int(stats['count']):6d} {_format_bytes(stats['total']):>12} "
+        f"{_format_bytes(stats['mean']):>12} {_format_bytes(stats['median']):>12} "
+        f"{_format_bytes(stats['min']):>12} {_format_bytes(stats['max']):>12} "
+        f"{_format_bytes(stats['p95']):>12} {_format_bytes(stats['std']):>12}"
+    )
+
+
+def reset_memory_trackers() -> None:
+    """Clear all recorded memory tracking data."""
+    MEMORY_REGISTRY.clear()
+
+
+__all__ = [
+    "MEMORY_REGISTRY",
+    "add_memory_tracker",
+    "compute_aggregate_memory_stats",
+    "save_memory_stats_csv",
+    "print_aggregate_memory_stats",
+    "compute_global_memory_stats",
+    "print_global_memory_stats",
+    "reset_memory_trackers",
+]
+
+
+if __name__ == "__main__":  # Simple demonstration
+
+    class Demo:
+        def __init__(self, device: str | None = None):
+            self.device = device or ("cuda" if _cuda_available() else "cpu")
+
+        def allocate(self, n: int = 1_000_000) -> int:
+            if not _cuda_available():
+                # Fallback: just create a CPU tensor
+                _ = [0] * n  # noqa: F841
+                return n
+            import torch  # local import to avoid mypy confusion
+
+            t = torch.empty(n, dtype=torch.float32, device=self.device)
+            # Perform an op to ensure allocation
+            t.uniform_()  # noqa: F841
+            return t.numel()
+
+        def noalloc(self):  # method with negligible allocation
+            return 42
+
+    demo = Demo()
+    add_memory_tracker(demo.allocate, "alloc")
+    add_memory_tracker(demo.noalloc, "noalloc")
+    add_memory_tracker(demo.allocate, "alloc")  # idempotent
+
+    for _ in range(5):
+        demo.allocate(200_000)
+        demo.noalloc()
+
+    print("\nAll memory stats:\n")
+    print_aggregate_memory_stats()
+
+    print("\nSingle (alloc):\n")
+    print_aggregate_memory_stats("alloc")
+
+    print("\nGlobal memory stats:\n")
+    print_global_memory_stats()

src/ctx_to_lora/tracker/timer.py

@@ -0,0 +1,332 @@
+"""Lightweight timing utilities for attaching runtime measurement to object methods.
+
+Usage:
+    x = SomeClass()
+    add_timer(x.some_method, "some_method")  # wraps the bound method in-place
+    x.some_method(...)
+    print_aggregate_timer_stats("some_method")
+
+Design notes:
+- add_timer mutates the instance by replacing the bound method with a timing wrapper.
+- Multiple calls to add_timer on the same (already wrapped) method are ignored to avoid double timing.
+- Global registry: { name: [float, ...] } storing individual call durations.
+- print_aggregate_timer_stats prints summary stats (count, total, mean, median, min, max, p95, last).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from statistics import mean, median, stdev
+from time import perf_counter
+from typing import Any
+
+# Global timer registry: name -> list of durations (seconds)
+TIMER_REGISTRY: dict[str, list[float]] = {}
+
+
+class _TimerWrapperMarker:
+    """Mixin marker to identify already wrapped callables."""
+
+    __slots__ = ("__wrapped_name__",)
+
+    def __init__(self, wrapped_name: str):
+        self.__wrapped_name__ = wrapped_name
+
+
+def add_timer(func: Callable, name: str) -> None:
+    """Attach a timing wrapper to a bound method.
+
+    Parameters
+    ----------
+    func : Callable
+        A *bound* method (e.g., instance.method). If an unbound function is provided
+        it will raise a ValueError (explicitness keeps behavior predictable).
+    name : str
+        Key under which durations are recorded in TIMER_REGISTRY.
+    """
+    # Basic validation (permit already wrapped functions for idempotency even though
+    # they sit as plain functions on the instance dict and thus lack __self__).
+    if not hasattr(func, "__self__") or getattr(func, "__self__") is None:
+        if getattr(func, "__is_timer_wrapper__", False):  # already wrapped, no-op
+            return
+        raise ValueError("add_timer expects a bound method: call with instance.method")
+
+    instance = func.__self__  # The object instance
+    method_name = getattr(func, "__name__", None)
+    if method_name is None:
+        raise ValueError("Cannot determine method name for provided callable")
+
+    # Prevent double-wrapping (idempotent behavior)
+    existing = getattr(instance, method_name, None)
+    if getattr(existing, "__is_timer_wrapper__", False):  # Already wrapped
+        return
+
+    orig_bound = func  # capture original bound method
+
+    def timed(*args: Any, **kwargs: Any):  # noqa: D401 - simple wrapper
+        start = perf_counter()
+        try:
+            return orig_bound(*args, **kwargs)
+        finally:
+            elapsed = perf_counter() - start
+            TIMER_REGISTRY.setdefault(name, []).append(elapsed)
+
+    # Mark wrapper to avoid double wrapping; preserve introspection hints.
+    timed.__name__ = method_name
+    timed.__doc__ = getattr(orig_bound, "__doc__")
+    timed.__qualname__ = getattr(orig_bound, "__qualname__", method_name)
+    timed.__is_timer_wrapper__ = True  # type: ignore[attr-defined]
+    timed.__wrapped__ = orig_bound  # type: ignore[attr-defined]
+    timed.__timer_name__ = name  # type: ignore[attr-defined]
+
+    setattr(instance, method_name, timed)
+
+
+def _format_seconds(sec: float) -> str:
+    if sec >= 1:
+        return f"{sec:8.3f}s"
+    if sec >= 1e-3:
+        return f"{sec * 1e3:8.3f}ms"
+    if sec >= 1e-6:
+        return f"{sec * 1e6:8.3f}µs"
+    return f"{sec * 1e9:8.3f}ns"
+
+
+def compute_aggregate_timer_stats(
+    name: str | None = None,
+) -> dict[str, dict[str, float]] | None:
+    """Compute aggregate timing statistics for specific timers.
+
+    Parameters
+    ----------
+    name : Optional[str]
+        Specific timer name to compute stats for. If None, all timers are computed.
+
+    Returns
+    -------
+    Optional[Dict[str, Dict[str, float]]]
+        None if no data, else dict mapping timer names to their stats.
+        Each stats dict contains: count, total, mean, median, min, max, p95, last, std.
+    """
+    if not TIMER_REGISTRY:
+        return None
+
+    keys = [name] if name else sorted(TIMER_REGISTRY.keys())
+    valid_keys = [k for k in keys if k in TIMER_REGISTRY and TIMER_REGISTRY[k]]
+    if not valid_keys:
+        return None
+
+    result = {}
+    for k in valid_keys:
+        data = TIMER_REGISTRY[k]
+        data_sorted = sorted(data)
+        cnt = len(data)
+        total = sum(data)
+        avg = mean(data)
+        med = median(data)
+        std = stdev(data) if cnt > 1 else 0.0
+        mn = data_sorted[0]
+        mx = data_sorted[-1]
+        p95_index = int(0.95 * (cnt - 1))
+        p95 = data_sorted[p95_index]
+        last = data[-1]
+        result[k] = {
+            "count": float(cnt),
+            "total": total,
+            "mean": avg,
+            "median": med,
+            "min": mn,
+            "max": mx,
+            "p95": p95,
+            "last": last,
+            "std": std,
+        }
+    return result
+
+
+def save_timer_stats_csv(file_path: str, name: str | None = None) -> None:
+    """Save aggregate timing statistics to a CSV file.
+
+    Parameters
+    ----------
+    file_path : str
+        Path where the CSV file will be saved.
+    name : Optional[str]
+        Specific timer name to export. If None, all timers are exported.
+    """
+    import csv
+
+    stats = compute_aggregate_timer_stats(name)
+    if stats is None:
+        raise ValueError("No timing data available to export")
+
+    with open(file_path, "w", newline="") as csvfile:
+        fieldnames = [
+            "name",
+            "count",
+            "total",
+            "mean",
+            "median",
+            "min",
+            "max",
+            "p95",
+            "last",
+            "std",
+        ]
+        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+        writer.writeheader()
+
+        for timer_name, data in stats.items():
+            row = {"name": timer_name}
+            row.update(data)
+            writer.writerow(row)
+
+
+def print_aggregate_timer_stats(name: str | None = None) -> None:
+    """Print aggregate timing statistics.
+
+    Parameters
+    ----------
+    name : Optional[str]
+        Specific timer name to report. If None, all timers are reported.
+    """
+    stats = compute_aggregate_timer_stats(name)
+    if stats is None:
+        print("[timer] No timing data collected.")
+        return
+
+    keys = [name] if name else sorted(TIMER_REGISTRY.keys())
+    missing = [k for k in keys if k not in TIMER_REGISTRY or not TIMER_REGISTRY[k]]
+    if missing:
+        print(f"[timer] No data for: {', '.join(missing)}")
+
+    if not stats:
+        return
+
+    header = (
+        f"{'name':20} {'count':>6} {'total':>10} {'mean':>10} {'median':>10} "
+        f"{'min':>10} {'max':>10} {'p95':>10} {'std':>10} {'last':>10}"
+    )
+    print(header)
+    print("-" * len(header))
+
+    for k, data in stats.items():
+        print(
+            f"{k:20} {int(data['count']):6d} {_format_seconds(data['total']):>10} "
+            f"{_format_seconds(data['mean']):>10} {_format_seconds(data['median']):>10} "
+            f"{_format_seconds(data['min']):>10} {_format_seconds(data['max']):>10} "
+            f"{_format_seconds(data['p95']):>10} {_format_seconds(data['std']):>10} "
+            f"{_format_seconds(data['last']):>10}"
+        )
+
+
+def compute_global_timer_stats() -> dict[str, float] | None:
+    """Compute aggregate statistics across all recorded timer values.
+
+    Returns
+    -------
+    Optional[Dict[str, float]]
+        None if no data recorded, else a dict with keys:
+        count, total, mean, median, min, max, p95, std.
+    """
+    if not TIMER_REGISTRY:
+        return None
+    all_values: list[float] = []
+    for lst in TIMER_REGISTRY.values():
+        all_values.extend(lst)
+    if not all_values:
+        return None
+    data_sorted = sorted(all_values)
+    cnt = len(all_values)
+    total = sum(all_values)
+    avg = mean(all_values)
+    med = median(all_values)
+    std = stdev(all_values) if cnt > 1 else 0.0
+    mn = data_sorted[0]
+    mx = data_sorted[-1]
+    p95_index = int(0.95 * (cnt - 1))
+    p95 = data_sorted[p95_index]
+    return {
+        "count": float(cnt),  # keep uniform numeric type
+        "total": total,
+        "mean": avg,
+        "median": med,
+        "min": mn,
+        "max": mx,
+        "p95": p95,
+        "std": std,
+    }
+
+
+def print_global_timer_stats() -> None:
+    """Pretty-print global aggregate stats across all timer entries."""
+    stats = compute_global_timer_stats()
+    if stats is None:
+        print("[timer] No timing data collected.")
+        return
+    header = (
+        f"{'scope':20} {'count':>6} {'total':>10} {'mean':>10} {'median':>10} "
+        f"{'min':>10} {'max':>10} {'p95':>10} {'std':>10}"
+    )
+    print(header)
+    print("-" * len(header))
+    print(
+        f"{'<ALL>':20} {int(stats['count']):6d} {_format_seconds(stats['total']):>10} "
+        f"{_format_seconds(stats['mean']):>10} {_format_seconds(stats['median']):>10} "
+        f"{_format_seconds(stats['min']):>10} {_format_seconds(stats['max']):>10} "
+        f"{_format_seconds(stats['p95']):>10} {_format_seconds(stats['std']):>10}"
+    )
+
+
+def reset_timers() -> None:
+    """Reset (clear) all recorded timing data."""
+    TIMER_REGISTRY.clear()
+
+
+__all__ = [
+    "TIMER_REGISTRY",
+    "add_timer",
+    "compute_aggregate_timer_stats",
+    "save_timer_stats_csv",
+    "print_aggregate_timer_stats",
+    "compute_global_timer_stats",
+    "print_global_timer_stats",
+    "reset_timers",
+]
+
+if __name__ == "__main__":
+    # Example usage / simple self-test.
+    import random
+    import time
+
+    class Demo:
+        def f1(self, n: int = 20_000) -> int:
+            # CPU-bound work
+            s = 0
+            for i in range(n):
+                s += i * i
+            return s
+
+        def f2(self) -> None:
+            # Simulate I/O or waiting
+            time.sleep(random.uniform(0.001, 0.003))
+
+    demo = Demo()
+
+    # Attach timers (idempotent: calling again does nothing harmful)
+    add_timer(demo.f1, "f1")
+    add_timer(demo.f2, "f2")
+    add_timer(demo.f1, "f1")  # demonstrate double-wrap prevention
+
+    for _ in range(5):
+        demo.f1(10_000)
+        demo.f2()
+
+    print("\nAll timers:\n")
+    print_aggregate_timer_stats()
+
+    print("\nSingle timer (f1):\n")
+    print_aggregate_timer_stats("f1")
+
+    print("\nGlobal aggregated stats:\n")
+    print_global_timer_stats()

src/ctx_to_lora/tracker/tracker.py

@@ -0,0 +1,303 @@
+"""Unified tracking interface combining timing and CUDA memory usage.
+
+Primary API
+-----------
+add_tracker(bound_method, name)
+    Wraps a bound instance method so that each invocation records:
+      - wall-clock duration (seconds) in timer.TIMER_REGISTRY[name]
+      - CUDA peak memory increase (bytes) in cuda_memory_tracker.MEMORY_REGISTRY[name]
+        (only if CUDA + torch available; otherwise memory list may stay empty / absent)
+
+print_tracker_stats(name=None)
+    Convenience printer that delegates to time + memory aggregate printers.
+
+Design
+------
+We implement a single wrapper (instead of nesting the individual timer & memory
+wrappers) to avoid multiple layers of indirection and to ensure the measured
+CUDA memory footprint reflects only the original method's body (excluding the
+separate timing wrapper's slight overhead). The wrapper is idempotent: repeated
+calls to add_tracker on the same method are ignored.
+
+This file depends on sibling modules:
+- tracker.timer
+- tracker.cuda_memory_tracker
+
+Both registries remain the single source of truth; no additional registry is introduced.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from time import perf_counter
+from typing import Any
+
+# Support both package (relative) and direct script execution.
+try:  # Package / normal import path
+    from .cuda_memory_tracker import (  # type: ignore
+        MEMORY_REGISTRY,
+        compute_aggregate_memory_stats,
+        print_aggregate_memory_stats,
+        print_global_memory_stats,
+        reset_memory_trackers,
+        save_memory_stats_csv,
+    )
+    from .timer import (  # type: ignore
+        TIMER_REGISTRY,
+        compute_aggregate_timer_stats,
+        print_aggregate_timer_stats,
+        print_global_timer_stats,
+        reset_timers,
+        save_timer_stats_csv,
+    )
+except Exception:  # pragma: no cover - fallback when executed directly
+    import pathlib
+    import sys
+
+    _this_file = pathlib.Path(__file__).resolve()
+    # project root is two levels up from tracker/ (i.e., .../src)
+    _src_root = _this_file.parents[2]
+    if str(_src_root) not in sys.path:
+        sys.path.insert(0, str(_src_root))
+    try:
+        from ctx_to_lora.tracker.cuda_memory_tracker import (  # type: ignore
+            MEMORY_REGISTRY,
+            compute_aggregate_memory_stats,
+            print_aggregate_memory_stats,
+            print_global_memory_stats,
+            reset_memory_trackers,
+            save_memory_stats_csv,
+        )
+        from ctx_to_lora.tracker.timer import (  # type: ignore
+            TIMER_REGISTRY,
+            compute_aggregate_timer_stats,
+            print_aggregate_timer_stats,
+            print_global_timer_stats,
+            reset_timers,
+            save_timer_stats_csv,
+        )
+    except Exception as e:  # If still failing, raise a clearer error.
+        raise ImportError(
+            f"Failed to import tracking dependencies; ensure project root on PYTHONPATH. Original: {e}"
+        )
+
+try:  # Optional torch import (lazy fallback if unavailable)
+    import torch  # type: ignore
+except Exception:  # pragma: no cover - torch absence path
+    torch = None  # type: ignore
+
+__all__ = [
+    "add_tracker",
+    "compute_tracker_stats",
+    "save_tracker_stats_csv",
+    "print_tracker_stats",
+    "print_global_tracker_stats",
+    "reset_trackers",
+]
+
+
+def _cuda_available() -> bool:
+    return bool(torch is not None and torch.cuda.is_available())
+
+
+def add_tracker(func: Callable, name: str) -> None:
+    """Attach a combined time + CUDA memory tracking wrapper to a bound method.
+
+    Parameters
+    ----------
+    func : Callable
+        A bound instance method (instance.method). Raises ValueError if unbound.
+    name : str
+        Registry key under which metrics are stored.
+    """
+    if not hasattr(func, "__self__") or getattr(func, "__self__") is None:
+        # Permit idempotent re-calls if already wrapped.
+        if getattr(func, "__is_tracker_wrapper__", False):
+            return
+        raise ValueError(
+            "add_tracker expects a bound method: call with instance.method"
+        )
+
+    instance = func.__self__  # underlying object
+    method_name = getattr(func, "__name__", None)
+    if method_name is None:
+        raise ValueError("Cannot determine method name for provided callable")
+
+    existing = getattr(instance, method_name, None)
+    if getattr(
+        existing, "__is_tracker_wrapper__", False
+    ):  # Already wrapped via unified tracker
+        return
+
+    # If already individually wrapped by timer or memory tracker, we still wrap only once more;
+    # future calls to add_tracker will become no-ops.
+    orig_bound = existing if existing is not None else func
+
+    def tracked(*args: Any, **kwargs: Any):  # noqa: D401 - combined wrapper
+        use_cuda = _cuda_available()
+        if use_cuda:
+            torch.cuda.synchronize()
+            torch.cuda.reset_peak_memory_stats()
+            start_alloc = torch.cuda.memory_allocated()
+        start_time = perf_counter()
+        try:
+            return orig_bound(*args, **kwargs)
+        finally:
+            elapsed = perf_counter() - start_time
+            TIMER_REGISTRY.setdefault(name, []).append(elapsed)
+            if use_cuda:
+                torch.cuda.synchronize()
+                peak_alloc = torch.cuda.max_memory_allocated()
+                peak_increase = peak_alloc - start_alloc
+                if peak_increase < 0:  # Safety guard (should not happen)
+                    peak_increase = 0
+                MEMORY_REGISTRY.setdefault(name, []).append(int(peak_increase))
+
+    # Introspection / idempotency markers
+    tracked.__name__ = method_name
+    tracked.__qualname__ = getattr(orig_bound, "__qualname__", method_name)
+    tracked.__doc__ = getattr(orig_bound, "__doc__")
+    tracked.__wrapped__ = orig_bound  # type: ignore[attr-defined]
+    tracked.__is_tracker_wrapper__ = True  # type: ignore[attr-defined]
+    tracked.__is_timer_wrapper__ = True  # type: ignore[attr-defined]
+    tracked.__is_memory_wrapper__ = True  # type: ignore[attr-defined]
+    tracked.__tracker_name__ = name  # type: ignore[attr-defined]
+
+    setattr(instance, method_name, tracked)
+
+
+def compute_tracker_stats(
+    name: str | None = None,
+) -> dict[str, dict[str, Any]] | None:
+    """Compute both timing and memory stats for a given name (or all if None).
+
+    Parameters
+    ----------
+    name : Optional[str]
+        Specific tracker name; if None, computes all.
+
+    Returns
+    -------
+    Optional[Dict[str, Dict[str, Any]]]
+        None if no data, else dict with 'timing' and 'memory' keys containing
+        their respective aggregate statistics.
+    """
+    timer_stats = compute_aggregate_timer_stats(name)
+    memory_stats = compute_aggregate_memory_stats(name)
+
+    if timer_stats is None and memory_stats is None:
+        return None
+
+    return {
+        "timing": timer_stats or {},
+        "memory": memory_stats or {},
+    }
+
+
+def save_tracker_stats_csv(file_path: str, name: str | None = None) -> None:
+    """Save both timing and memory stats to separate CSV files.
+
+    Parameters
+    ----------
+    file_path : str
+        Base path for CSV files. Will create file_path_timing.csv and file_path_memory.csv
+    name : Optional[str]
+        Specific tracker name to export. If None, all trackers are exported.
+    """
+    import os
+
+    os.makedirs(os.path.dirname(file_path), exist_ok=True)
+
+    base_path = os.path.splitext(file_path)[0]
+    timer_path = f"{base_path}_timing.csv"
+    memory_path = f"{base_path}_memory.csv"
+
+    # Save timing stats if available
+    timer_stats = compute_aggregate_timer_stats(name)
+    if timer_stats is not None:
+        save_timer_stats_csv(timer_path, name)
+
+    # Save memory stats if available
+    memory_stats = compute_aggregate_memory_stats(name)
+    if memory_stats is not None:
+        save_memory_stats_csv(memory_path, name)
+
+    # If no data at all, raise an error
+    if timer_stats is None and memory_stats is None:
+        raise ValueError("No tracking data available to export")
+
+
+def print_tracker_stats(name: str | None = None) -> None:
+    """Print both timing and memory stats for a given name (or all if None).
+
+    Parameters
+    ----------
+    name : Optional[str]
+        Specific tracker name; if None, prints all.
+    """
+    print("[tracker] Timing stats:")
+    print_aggregate_timer_stats(name)
+    print("\n[tracker] CUDA memory stats:")
+    print_aggregate_memory_stats(name)
+
+
+def print_global_tracker_stats() -> None:
+    """Print global aggregate timing and memory stats."""
+    print("[tracker] Global timing stats:")
+    print_global_timer_stats()
+    print("\n[tracker] Global CUDA memory stats:")
+    print_global_memory_stats()
+
+
+def reset_trackers() -> None:
+    """Reset all timer and memory tracking data."""
+    reset_timers()
+    reset_memory_trackers()
+
+
+if __name__ == "__main__":  # Demonstration
+    import random
+    import time
+
+    class Demo:
+        def compute(self, n: int = 25_000) -> int:
+            # CPU-bound work
+            s = 0
+            for i in range(n):
+                s += i * i
+            return s
+
+        def gpu_alloc(self, n: int = 500_000):
+            if not _cuda_available():
+                # Simulate light wait to differentiate timing
+                time.sleep(random.uniform(0.01, 0.05))
+                return None
+            t = torch.empty(n, dtype=torch.float32, device="cuda")
+            t.uniform_()  # ensure usage
+            return t.sum().item()
+
+    demo = Demo()
+
+    add_tracker(demo.compute, "compute")
+    add_tracker(demo.gpu_alloc, "gpu_alloc")
+    # Idempotent re-call
+    add_tracker(demo.compute, "compute")
+
+    for _ in range(5):
+        demo.compute(15_000)
+        demo.gpu_alloc(300_000)
+
+    print_tracker_stats()
+    print("\n--- Global Combined Stats ---\n")
+    print_global_tracker_stats()
+
+    # Demonstrate CSV export
+    print("\n[tracker] Saving stats to CSV files...\n")
+    csv_path = "/tmp/tracker_demo_stats.csv"
+    save_tracker_stats_csv(csv_path)
+    print(f"Exported timing stats to: {csv_path.replace('.csv', '_timing.csv')}")
+    print(f"Exported memory stats to: {csv_path.replace('.csv', '_memory.csv')}")
+
+    print("\n[tracker] Resetting registries...\n")
+    reset_trackers()
+    print_tracker_stats()

train.py

@@ -302,6 +302,11 @@ def main():
                 tokenized_ds[split][ds_name] = ds
 
     train_ds = tokenized_ds["train"]
+    if data_args.max_train_samples_per_ds is not None:
+        for ds_name, ds in train_ds.items():
+            if data_args.max_train_samples_per_ds >= len(ds):
+                continue
+            train_ds[ds_name] = ds.take(data_args.max_train_samples_per_ds)
     logging.info(f"train_ds: {train_ds}")
 
     val_ds = dict()