iclr 2026 (tmp 2)

tmp/icl_results.py

@@ -0,0 +1,204 @@
+import json
+from pathlib import Path
+
+all_results_path = Path(
+    "eval_results/google/gemma-2-2b-it/20250918-111058_734998c4/all_results.json"
+)
+all_result_modulated_path = Path(
+    "../train_outputs/runs/Aug10_08-29-46_sakura-003_2000_f73050aa/eval-results-100000/all_results.json"
+)
+with open(all_results_path) as f:
+    all_results = json.load(f)
+
+with open(all_result_modulated_path) as f:
+    all_results_modulated = json.load(f)
+
+# Find keys with test_metaicl_macro_f1 and test_metaicl_no_context_macro_f1 prefixes
+metaicl_keys = [k for k in all_results.keys() if k.startswith("test_metaicl_macro_f1")]
+no_context_keys = [
+    k for k in all_results.keys() if k.startswith("test_metaicl_no_context_macro_f1")
+]
+
+# Sort keys for consistent ordering
+metaicl_keys.sort()
+no_context_keys.sort()
+
+print("Comparison of MetaICL macro F1 scores:")
+print("=" * 140)
+print(f"{'Without Context':<40} {'With Context':<40} {'Modulated (No Context)':<40}")
+print("-" * 140)
+
+# Create a mapping of task names to values for easier comparison
+metaicl_dict = {}
+no_context_dict = {}
+modulated_dict = {}
+
+for key in metaicl_keys:
+    task_name = key.replace("test_metaicl_macro_f1_", "")
+    metaicl_dict[task_name] = all_results[key]
+
+for key in no_context_keys:
+    task_name = key.replace("test_metaicl_no_context_macro_f1_", "")
+    no_context_dict[task_name] = all_results[key]
+
+# Get modulated results (assuming they use the same key pattern as with context)
+for key in all_results_modulated.keys():
+    if key.startswith("test_metaicl_macro_f1_"):
+        task_name = key.replace("test_metaicl_macro_f1_", "")
+        modulated_dict[task_name] = all_results_modulated[key]
+
+# Get all unique task names from the base results (this determines ordering)
+all_tasks = set(metaicl_dict.keys()) | set(no_context_dict.keys())
+
+# Create list of tasks with their differences for sorting (based on without context vs with context)
+task_data = []
+for task in all_tasks:
+    with_context = metaicl_dict.get(task, "N/A")
+    without_context = no_context_dict.get(task, "N/A")
+    modulated = modulated_dict.get(task, "N/A")
+
+    if isinstance(with_context, (int, float)) and isinstance(
+        without_context, (int, float)
+    ):
+        diff_with_context = with_context - without_context
+        modulated_diff = (
+            modulated - without_context
+            if isinstance(modulated, (int, float))
+            else float("-inf")
+        )
+        task_data.append(
+            (
+                task,
+                without_context,
+                with_context,
+                modulated,
+                diff_with_context,
+                modulated_diff,
+            )
+        )
+    else:
+        # Put tasks with missing data at the end
+        task_data.append(
+            (
+                task,
+                without_context,
+                with_context,
+                modulated,
+                float("-inf"),
+                float("-inf"),
+            )
+        )
+
+# Sort by difference (with context - without context) from high to low
+task_data.sort(key=lambda x: x[4], reverse=True)
+
+for (
+    task,
+    without_context,
+    with_context,
+    modulated,
+    diff_with_context,
+    modulated_diff,
+) in task_data:
+    if isinstance(with_context, (int, float)) and isinstance(
+        without_context, (int, float)
+    ):
+        modulated_str = (
+            f"{modulated:.4f}"
+            if isinstance(modulated, (int, float))
+            else str(modulated)
+        )
+        modulated_diff_str = (
+            f"(+{modulated_diff:+.4f})"
+            if isinstance(modulated_diff, (int, float))
+            and modulated_diff != float("-inf")
+            else ""
+        )
+        print(
+            f"{task}: {without_context:.4f} {task}: {with_context:.4f} (+{diff_with_context:+.4f}) {task}: {modulated_str} {modulated_diff_str}"
+        )
+    else:
+        print(f"{task}: {without_context} {task}: {with_context} {task}: {modulated}")
+
+IMPROVEMENT_THRESHOLD = 0.10
+# Filter tasks where with_context shows >5% increase over without_context
+high_improvement_tasks = []
+for (
+    task,
+    without_context,
+    with_context,
+    modulated,
+    diff_with_context,
+    modulated_diff,
+) in task_data:
+    if (
+        isinstance(with_context, (int, float))
+        and isinstance(without_context, (int, float))
+        and isinstance(diff_with_context, (int, float))
+        and diff_with_context > IMPROVEMENT_THRESHOLD
+    ):
+        high_improvement_tasks.append(
+            (
+                task,
+                without_context,
+                with_context,
+                modulated,
+                diff_with_context,
+                modulated_diff,
+            )
+        )
+
+if high_improvement_tasks:
+    print("\n" + "=" * 80)
+    print(
+        f"Tasks with >{IMPROVEMENT_THRESHOLD * 100}% improvement when adding context ({len(high_improvement_tasks)} tasks):"
+    )
+    print("=" * 80)
+
+    # Calculate averages
+    with_context_improvements = [diff for _, _, _, _, diff, _ in high_improvement_tasks]
+    modulated_improvements = [
+        mod_diff
+        for _, _, _, _, _, mod_diff in high_improvement_tasks
+        if isinstance(mod_diff, (int, float)) and mod_diff != float("-inf")
+    ]
+
+    avg_with_context_improvement = sum(with_context_improvements) / len(
+        with_context_improvements
+    )
+    avg_modulated_improvement = (
+        sum(modulated_improvements) / len(modulated_improvements)
+        if modulated_improvements
+        else 0
+    )
+
+    print(f"Average improvement with context: +{avg_with_context_improvement:.4f}")
+    print(f"Average improvement with modulated: +{avg_modulated_improvement:.4f}")
+
+    print("\nDetailed breakdown:")
+    for (
+        task,
+        without_context,
+        with_context,
+        modulated,
+        diff_with_context,
+        modulated_diff,
+    ) in high_improvement_tasks:
+        modulated_str = (
+            f"{modulated:.4f}"
+            if isinstance(modulated, (int, float))
+            else str(modulated)
+        )
+        modulated_diff_str = (
+            f"(+{modulated_diff:+.4f})"
+            if isinstance(modulated_diff, (int, float))
+            and modulated_diff != float("-inf")
+            else ""
+        )
+        print(
+            f"  {task}: {without_context:.4f} → {with_context:.4f} (+{diff_with_context:+.4f}) | Modulated: {modulated_str} {modulated_diff_str}"
+        )
+else:
+    print(
+        f"\nNo tasks found with >{IMPROVEMENT_THRESHOLD * 100}% improvement when adding context."
+    )

tmp/test_chat.py

@@ -0,0 +1,105 @@
+import os
+
+import torch
+
+from ctx_to_lora.model_loading import get_model, get_tokenizer
+
+# tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
+# model = AutoModelForCausalLM.from_pretrained(
+#     "google/gemma-2-2b-it",
+#     device_map="auto",
+#     torch_dtype=torch.bfloat16,
+#     attn_implementation="sdpa",
+# )
+
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
+os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"
+
+# torch.use_deterministic_algorithms(True, warn_only=True)
+torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
+torch.backends.cudnn.benchmark = False
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cudnn.allow_tf32 = False
+
+# model, tokenizer = get_model_and_tokenizer(
+#     "google/gemma-3-1b-it",
+#     train=False,
+#     requires_grad=False,
+#     use_flash_attn=True,
+#     peft_config=None,
+#     model_kwargs={"attn_implementation": "flash_attention_2"},
+#     tokenizer_kwargs=None,
+#     device="cuda",
+#     dtype=torch.bfloat16,
+# )
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_kwargs = {"attn_implementation": "flash_attention_2"}
+# model_name_or_path = "Qwen/Qwen3-1.7B"
+model_name_or_path = "google/gemma-2-2b-it"
+# model_name_or_path = "google/gemma-3-1b-it"
+tokenizer = get_tokenizer(model_name_or_path, train=False)
+model = get_model(
+    model_name_or_path,
+    train=False,
+    requires_grad=False,
+    model_kwargs=model_kwargs,
+    use_flash_attn=True,
+    dtype=torch.bfloat16,
+)
+
+qa = (
+    "Architecturally, the school has a Catholic character. "
+    "Atop the Main Building's gold dome is a golden statue of the Virgin Mary. "
+    "Immediately in front of the Main Building and facing it, "
+    'is a copper statue of Christ with arms upraised with the legend "Venite Ad Me Omnes". '
+    "Next to the Main Building is the Basilica of the Sacred Heart. "
+    "Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. "
+    "It is a replica of the grotto at Lourdes, "
+    "France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous "
+    "in 1858. At the end of the main drive (and in a direct line that connects through 3 "
+    "statues and the Gold Dome), is a simple, modern stone statue of Mary."
+    "\n\nAnswer the following question. "
+    "Output only the answer and do not output any other words."
+    "\n\nQuestion: To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?"
+)
+# x = " | ".join([str(random.randint(0, 1000)) for _ in range(800)])
+# qa = (
+#     " | ".join([str(random.randint(0, 1000)) for _ in range(10000)])
+#     + "\n\nRepeat the text above."
+# )
+
+messages = [
+    [{"role": "system", "content": ""}, {"role": "user", "content": qa}],
+    [{"role": "system", "content": ""}, {"role": "user", "content": "Hello."}],
+]
+chat_ids = tokenizer.apply_chat_template(
+    messages,
+    tokenize=True,
+    padding=True,
+    truncation=False,
+    add_special_token=False,
+    add_generation_prompt=True,
+    return_tensors="pt",
+    return_dict=True,
+).to(device)
+print(chat_ids)
+print(f"{len(chat_ids['input_ids'][0])} tokens in input")
+print(tokenizer.decode(chat_ids["input_ids"][0]))
+
+
+print(model(**chat_ids).logits)
+
+outputs = model.generate(
+    **chat_ids,
+    max_new_tokens=2**12,
+    do_sample=False,
+    temperature=0.0,
+)
+for output in outputs:
+    print(tokenizer.decode(output))
+    print("-" * 100)
+
+print(outputs)
+breakpoint()

tmp/test_detokenize.py

@@ -0,0 +1,35 @@
+from ctx_to_lora.model_loading import get_tokenizer
+
+model_name_or_path = "google/gemma-2-2b-it"
+tokenizer = get_tokenizer(model_name_or_path, train=True)
+
+
+qa = (
+    "Architecturally, the school has a Catholic character. "
+    "Atop the Main Building's gold dome is a golden statue of the Virgin Mary. "
+    "Immediately in front of the Main Building and facing it, "
+    'is a copper statue of Christ with arms upraised with the legend "Venite Ad Me Omnes". '
+    "Next to the Main Building is the Basilica of the Sacred Heart. "
+    "Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. "
+    "It is a replica of the grotto at Lourdes, "
+    "France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous "
+    "in 1858. At the end of the main drive (and in a direct line that connects through 3 "
+    "statues and the Gold Dome), is a simple, modern stone statue of Mary."
+    "\n\nAnswer the following question. "
+    "Output only the answer and do not output any other words."
+    "\n\nQuestion: To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?"
+)
+
+
+messages = [{"role": "system", "content": ""}, {"role": "user", "content": qa}]
+chat_ids = tokenizer.apply_chat_template(
+    messages,
+    tokenize=True,
+    add_special_token=False,
+    add_generation_prompt=True,
+    return_tensors="pt",
+    return_dict=True,
+)
+print(chat_ids)
+print(tokenizer.batch_decode(chat_ids["input_ids"]))
+breakpoint()

tmp/test_emb.py

@@ -0,0 +1,51 @@
+import os
+
+import torch
+
+from ctx_to_lora.model_loading import get_model, get_tokenizer
+
+# tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
+# model = AutoModelForCausalLM.from_pretrained(
+#     "google/gemma-2-2b-it",
+#     device_map="auto",
+#     torch_dtype=torch.bfloat16,
+#     attn_implementation="sdpa",
+# )
+
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
+os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"
+
+# torch.use_deterministic_algorithms(True, warn_only=True)
+torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
+torch.backends.cudnn.benchmark = False
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cudnn.allow_tf32 = False
+
+# model, tokenizer = get_model_and_tokenizer(
+#     "google/gemma-3-1b-it",
+#     train=False,
+#     requires_grad=False,
+#     use_flash_attn=True,
+#     peft_config=None,
+#     model_kwargs={"attn_implementation": "flash_attention_2"},
+#     tokenizer_kwargs=None,
+#     device="cuda",
+#     dtype=torch.bfloat16,
+# )
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_kwargs = {"attn_implementation": "flash_attention_2"}
+# model_name_or_path = "google/gemma-2-2b-it"
+model_name_or_path = "google/gemma-2-2b-it"
+tokenizer = get_tokenizer(model_name_or_path, train=False)
+model = get_model(
+    model_name_or_path,
+    train=False,
+    requires_grad=False,
+    model_kwargs=model_kwargs,
+    use_flash_attn=True,
+    dtype=torch.bfloat16,
+)
+
+breakpoint()

tmp/test_logits.py

@@ -0,0 +1,121 @@
+import os
+
+import torch
+
+from ctx_to_lora.model_loading import get_model, get_tokenizer
+
+# tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
+# model = AutoModelForCausalLM.from_pretrained(
+#     "google/gemma-2-2b-it",
+#     device_map="auto",
+#     torch_dtype=torch.bfloat16,
+#     attn_implementation="sdpa",
+# )
+
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
+os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"
+
+# torch.use_deterministic_algorithms(True, warn_only=True)
+torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
+torch.backends.cudnn.benchmark = False
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cudnn.allow_tf32 = False
+
+# model, tokenizer = get_model_and_tokenizer(
+#     "google/gemma-3-1b-it",
+#     train=False,
+#     requires_grad=False,
+#     use_flash_attn=True,
+#     peft_config=None,
+#     model_kwargs={"attn_implementation": "flash_attention_2"},
+#     tokenizer_kwargs=None,
+#     device="cuda",
+#     dtype=torch.bfloat16,
+# )
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_kwargs = {"attn_implementation": "flash_attention_2"}
+# model_name_or_path = "google/gemma-2-2b-it"
+model_name_or_path = "google/gemma-3-1b-it"
+tokenizer = get_tokenizer(model_name_or_path, train=False)
+model = get_model(
+    model_name_or_path,
+    train=False,
+    requires_grad=False,
+    model_kwargs=model_kwargs,
+    use_flash_attn=True,
+    dtype=torch.bfloat16,
+)
+# model.config.pad_token_id = tokenizer.pad_token_id
+# if getattr(model, "generation_config", None):
+#     model.generation_config.pad_token_id = tokenizer.pad_token_id
+# print(model.name_or_path)
+
+input_text = [
+    "Write me a poem about Machine Learning.",
+    "Write me a poem about cats.",
+    "Who is Donal Trump?",
+    "Who is Mark Zuckerburg?",
+    "Tell me about Thailand.",
+    "Is Japan hot in the summer?",
+    "How computers work?",
+    "Who sells iPhone?",
+    "What exclusive games are on PS5?",
+    "Mechanical vs rubber dome keyboards",
+]
+input_ids = [tokenizer(t, return_tensors="pt").to(device) for t in input_text]
+
+padding_kwargs = dict(padding=True, padding_side="left", return_tensors="pt")
+input_ids = tokenizer.pad(
+    {k: [x[k][0] for x in input_ids] for k in input_ids[0].keys()},
+    **padding_kwargs,
+).to(device)
+print(input_ids)
+outputs = model.generate(**input_ids, do_sample=False)
+for output in outputs:
+    print(tokenizer.decode(output))
+    print("-" * 100)
+
+print(model(**input_ids).logits)
+
+# qa = (
+#     "Architecturally, the school has a Catholic character. "
+#     "Atop the Main Building's gold dome is a golden statue of the Virgin Mary. "
+#     "Immediately in front of the Main Building and facing it, "
+#     'is a copper statue of Christ with arms upraised with the legend "Venite Ad Me Omnes". '
+#     "Next to the Main Building is the Basilica of the Sacred Heart. "
+#     "Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. "
+#     "It is a replica of the grotto at Lourdes, "
+#     "France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous "
+#     "in 1858. At the end of the main drive (and in a direct line that connects through 3 "
+#     "statues and the Gold Dome), is a simple, modern stone statue of Mary."
+#     "\n\nAnswer the following question. "
+#     "Output only the answer and do not output any other words."
+#     "\n\nQuestion: To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?"
+# )
+
+
+# messages = [{"role": "system", "content": ""}, {"role": "user", "content": qa}]
+# chat_ids = tokenizer.apply_chat_template(
+#     messages,
+#     tokenize=True,
+#     add_special_token=False,
+#     add_generation_prompt=True,
+#     return_tensors="pt",
+#     return_dict=True,
+# ).to(device)
+# print(chat_ids)
+# print(tokenizer.decode(chat_ids["input_ids"][0]))
+# outputs = model.generate(
+#     **chat_ids,
+#     max_new_tokens=100,
+#     do_sample=False,
+#     temperature=0.0,
+# )
+
+# print(model(**chat_ids).logits)
+
+# for output in outputs:
+#     print(tokenizer.decode(output))
+#     print("-" * 100)

tmp/test_mem_alloc.py

@@ -0,0 +1,151 @@
+import os
+import time
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+def bytes_to_mib(x: int) -> float:
+    return x / (1024**2)
+
+
+def bytes_to_gib(x: int) -> float:
+    return x / (1024**3)
+
+
+def pick_dtype() -> torch.dtype:
+    if torch.cuda.is_available():
+        # Prefer bfloat16 if supported, else float16, else float32
+        if torch.cuda.is_bf16_supported():
+            return torch.bfloat16
+        return torch.float16
+    return torch.float32
+
+
+def main():
+    model_name = os.environ.get("MODEL_NAME", "google/gemma-2-2b-it")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = pick_dtype()
+
+    if device.type != "cuda":
+        print(
+            "CUDA is not available; this script is intended for GPU memory measurements."
+        )
+        return
+
+    torch.cuda.empty_cache()
+    torch.cuda.synchronize()
+
+    # Load tokenizer on CPU (doesn't affect CUDA mem)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    # Measure peak memory during model load
+    torch.cuda.reset_peak_memory_stats(device)
+    load_start = time.time()
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=dtype,
+        device_map=None,  # load on CPU first, then move to desired single device
+    )
+    model.to(device)
+    model.eval()
+
+    torch.cuda.synchronize()
+    model_load_peak = torch.cuda.max_memory_allocated(device)
+    load_elapsed = time.time() - load_start
+
+    # Baseline after load
+    baseline_alloc = torch.cuda.memory_allocated(device)
+
+    # Prepare a chat-style prompt if available; fallback to plain text
+    doc = """
+Matryoshka dolls (Russian: матрёшка, romanized: matryoshka/ˌmætriˈɒʃkə/), also known as stacking dolls, nesting dolls, Russian tea dolls, or Russian dolls,[1] are a set of wooden dolls of decreasing size placed one inside another. The name Matryoshka is a diminutive form of Matryosha (Матрёша), in turn a hypocorism of the Russian female first name Matryona (Матрёна).[2]
+
+A set of matryoshkas consists of a wooden figure, which separates at the middle, top from bottom, to reveal a smaller figure of the same sort inside, which has, in turn, another figure inside of it, and so on.
+
+The first Russian nested doll set was made in 1890 by woodturning craftsman and wood carver Vasily Zvyozdochkin from a design by Sergey Malyutin, who was a folk crafts painter at Abramtsevo. Traditionally the outer layer is a woman, dressed in a Russian sarafan dress. The figures inside may be of any gender; the smallest, innermost doll is typically a baby turned from a single piece of wood. Much of the artistry is in the painting of each doll, which can be very elaborate. The dolls often follow a theme; the themes may vary, from fairy tale characters to Soviet leaders. In some countries, matryoshka dolls are often referred to as babushka dolls, though they are not known by this name in Russian; babushka (бабушка) means 'grandmother; old woman'.[3]
+History
+The original matryoshka set by Zvyozdochkin and Malyutin, 1892
+
+The first Russian nested doll set was carved in 1890 at the Children's Education Workshop by Vasily Zvyozdochkin and designed by Sergey Malyutin, who was a folk crafts painter in the Abramtsevo estate of Savva Mamontov, a Russian industrialist and patron of arts.[4][5] Mamontov's brother, Anatoly Ivanovich Mamontov (1839–1905), created the Children's Education Workshop to make and sell children's toys. The doll set was painted by Malyutin. Malyutin's doll set consisted of eight dolls—the outermost was a mother in a traditional dress holding a red-combed rooster. The inner dolls were her children, girls and a boy, and the innermost a baby. The Children's Education Workshop was closed in the late 1890s, but the tradition of the matryoshka simply relocated to Sergiyev Posad, the Russian city known as a toy-making center since the fourteenth century.[6][4]
+
+The inspiration for matryoshka dolls is not clear. Matryoshka dolls may have been inspired by a nesting doll imported from Japan.[5][7] The Children's Education workshop where Zvyozdochkin was a lathe operator received a five piece, cylinder-shaped nesting doll featuring Fukuruma (Fukurokuju) in the late 1890s,[8] which is now part of the collection at the Sergiev Posad Museum of Toys.[8] Other east Asian dolls share similarities with matryoshka dolls such as the Kokeshi dolls,[4][9] originating in Northern Honshū, the main island of Japan, although they cannot be placed one inside another, and the round hollow daruma doll depicting a Buddhist monk.[9][10] Another possible source of inspiration is the nesting Easter eggs produced on a lathe by Russian woodworkers during the late 19th Century.[3][11]
+
+Savva Mamontov's wife presented a set of matryoshka dolls at the Exposition Universelle in Paris in 1900, and the toy earned a bronze medal. Soon after, matryoshka dolls were being made in several places in Russia and shipped around the world. 
+    """
+    messages = [
+        {
+            "role": "user",
+            "content": doc + "\n\nSummarize the article.",
+        },
+    ]
+
+    prompt_text = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    inputs = tokenizer(prompt_text, return_tensors="pt").to(device)
+
+    # Optional quick warmup to stabilize kernels (doesn't affect peak calc since we reset after)
+    with torch.no_grad():
+        _ = model.generate(**inputs, max_new_tokens=1)
+        torch.cuda.synchronize()
+
+    # Measure generation peak memory for 512 new tokens
+    torch.cuda.reset_peak_memory_stats(device)
+    gen_start = time.time()
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=int(os.environ.get("MAX_NEW_TOKENS", 512)),
+            do_sample=False,
+            temperature=None,
+            top_p=None,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+    torch.cuda.synchronize()
+    gen_elapsed = time.time() - gen_start
+
+    gen_peak_total = torch.cuda.max_memory_allocated(
+        device
+    )  # includes baseline present at reset
+    gen_additional = max(0, gen_peak_total - baseline_alloc)
+
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Print summary
+    print("=== Memory allocation summary (CUDA) ===")
+    print(f"Device: {device} | {torch.cuda.get_device_name(device)}")
+    print(f"Model: {model_name}")
+    print(f"DType: {dtype}")
+    print()
+    print("Model load:")
+    print(
+        f"  Peak allocated during load: {bytes_to_mib(model_load_peak):.2f} MiB ({bytes_to_gib(model_load_peak):.2f} GiB)"
+    )
+    print(f"  Time to load + move to device: {load_elapsed:.2f}s")
+    print(
+        f"  Baseline after load (current allocated): {bytes_to_mib(baseline_alloc):.2f} MiB ({bytes_to_gib(baseline_alloc):.2f} GiB)"
+    )
+    print()
+    print("Generation (max_new_tokens=512):")
+    print(
+        f"  Peak allocated total during generation: {bytes_to_mib(gen_peak_total):.2f} MiB ({bytes_to_gib(gen_peak_total):.2f} GiB)"
+    )
+    print(
+        f"  Additional peak over baseline during generation: {bytes_to_mib(gen_additional):.2f} MiB ({bytes_to_gib(gen_additional):.2f} GiB)"
+    )
+    print(f"  Time to generate: {gen_elapsed:.2f}s")
+    print()
+    # Response preview
+    preview = decoded
+    if len(preview) > 800:
+        preview = preview[:800] + "\n... [truncated]"
+    print("=== Response preview ===")
+    print(preview)
+
+
+if __name__ == "__main__":
+    main()

tmp/test_pad.py

@@ -0,0 +1,8 @@
+import torch
+from transformers import AutoTokenizer
+
+tk = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
+x = [torch.rand(5, 3), torch.rand(10, 3), torch.rand(7, 3)]
+padded_x = tk.pad({"input_ids": x}, padding=True)
+print(padded_x)
+breakpoint()

tmp/torch_profiler_test.py

@@ -0,0 +1,69 @@
+# (c) Meta Platforms, Inc. and affiliates.
+import logging
+import socket
+from datetime import datetime
+
+import torch
+from torch.autograd.profiler import record_function
+from torchvision import models
+
+logging.basicConfig(
+    format="%(levelname)s:%(asctime)s %(message)s",
+    level=logging.INFO,
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger: logging.Logger = logging.getLogger(__name__)
+logger.setLevel(level=logging.INFO)
+
+TIME_FORMAT_STR: str = "%b_%d_%H_%M_%S"
+
+
+def trace_handler(prof: torch.profiler.profile):
+    # Prefix for file names.
+    host_name = socket.gethostname()
+    timestamp = datetime.now().strftime(TIME_FORMAT_STR)
+    file_prefix = f"{host_name}_{timestamp}"
+
+    # Construct the trace file.
+    prof.export_chrome_trace(f"{file_prefix}.json.gz")
+
+    # Construct the memory timeline file.
+    prof.export_memory_timeline(f"{file_prefix}.html", device="cuda:0")
+
+
+def run_resnet50(num_iters=5, device="cuda:0"):
+    model = models.resnet50().to(device=device)
+    inputs = torch.randn(1, 3, 224, 224, device=device)
+    labels = torch.rand_like(model(inputs))
+    optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)
+    loss_fn = torch.nn.CrossEntropyLoss()
+
+    with torch.profiler.profile(
+        activities=[
+            torch.profiler.ProfilerActivity.CPU,
+            torch.profiler.ProfilerActivity.CUDA,
+        ],
+        schedule=torch.profiler.schedule(wait=0, warmup=0, active=6, repeat=1),
+        record_shapes=True,
+        profile_memory=True,
+        with_stack=True,
+        on_trace_ready=trace_handler,
+    ) as prof:
+        for _ in range(num_iters):
+            prof.step()
+            with record_function("## forward ##"):
+                pred = model(inputs)
+
+            with record_function("## backward ##"):
+                loss_fn(pred, labels).backward()
+
+            with record_function("## optimizer ##"):
+                optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+
+
+if __name__ == "__main__":
+    # Warm up
+    run_resnet50()
+    # Run the resnet50 model
+    run_resnet50()

tmp/visualize_ds_len.py

@@ -0,0 +1,101 @@
+import os
+from functools import partial
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from ctx_to_lora.data.processing import get_tokenized_dataset
+from ctx_to_lora.model_loading import get_tokenizer
+from ctx_to_lora.utils import check_is_iterable
+
+tokenizer_kwargs = {"max_length": 2**13}
+ctx_tokenizer_kwargs = {"max_length": 2**13}  # not used for now
+tokenizer = get_tokenizer("google/gemma-2-2b-it", train=True)
+ctx_tokenizer = get_tokenizer("google/gemma-2-2b-it", train=True)
+
+_get_tokenized_dataset = partial(
+    get_tokenized_dataset,
+    base_model_max_len=2**13,
+    tokenizer=tokenizer,
+    tokenizer_kwargs=tokenizer_kwargs,
+    ctx_model_max_len=2**13,
+    ctx_tokenizer=ctx_tokenizer,
+    ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
+    add_ctx_to_chat=False,
+    add_repeat_prompt=False,
+    add_negative_prompt=False,
+    use_kl_loss=False,
+    # streaming=data_args.streaming,
+)
+
+train_ds_names = [
+    # "fw_qa",
+    # "fw_qa_large",
+    # "ctx_qa",
+    # "pwc",
+    # "hotpot_qa",
+    # "squad",
+    # "drop",
+    # "narrativeqa",
+    # "quoref",
+    # "ropes",
+    # "synthetic_convqa",
+    # "booksum",
+    # "gov_report",
+    # "wikitext-2",
+    # "wikitext-103",
+    # "fw_qa_2"
+    # "fw_qa_xl",
+    # "fw_qa_3_medium",
+    "self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_0_tiny",
+    "self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_1_tiny",
+    "self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_2_tiny",
+    "self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_3_tiny",
+]
+tokenized_ds = dict()
+val_ds_names = []  # ["fw_qa_large", "ctx_qa", "pwc", "hotpot_qa", "squad"]
+for split, ds_names in zip(["train", "validation"], [train_ds_names, val_ds_names]):
+    tokenized_ds[split] = {
+        os.path.basename(ds_name): _get_tokenized_dataset(ds_name, split)
+        for ds_name in ds_names
+    }
+    for ds_name in tokenized_ds[split]:
+        print(ds_name)
+        print(f"Num samples: {len(tokenized_ds[split][ds_name])}")
+        sample = tokenized_ds[split][ds_name][0]
+
+        # TODO: let's investigate samples that are very long and very short
+        if check_is_iterable(sample["input_ids"]):
+            tokens = []
+            for sample in tokenized_ds[split][ds_name]:
+                tokens += [len(ids) for ids in sample["input_ids"]]
+        else:
+            tokens = [
+                len(sample["input_ids"]) for sample in tokenized_ds[split][ds_name]
+            ]
+
+        print(f"Num avg. tokens: {sum(tokens) / len(tokens)}")
+        print(f"Num max. tokens: {max(tokens)}")
+        print(f"Num min. tokens: {min(tokens)}")
+        print(f"Num std. tokens: {np.std(tokens)}")
+        plt.hist(tokens, bins=100)
+        plt.savefig(f"tmp/n_token_hist_{split}_{ds_name}.png")
+        plt.close()
+
+        long_indices = [i for i, n_tokens in enumerate(tokens) if n_tokens > 100]
+        short_indices = [i for i, n_tokens in enumerate(tokens) if n_tokens < 30]
+        breakpoint()
+        for i in long_indices:
+            print(f"Long sample {i}: {tokenized_ds[split][ds_name][i]}")
+        breakpoint()
+        for i in short_indices:
+            print(f"Short sample {i}: {tokenized_ds[split][ds_name][i]}")
+
+        ctx_tokens = [len(sample["ctx_ids"]) for sample in tokenized_ds[split][ds_name]]
+        print(f"Num avg. ctx tokens: {sum(ctx_tokens) / len(ctx_tokens)}")
+        print(f"Num max. ctx tokens: {max(ctx_tokens)}")
+        print(f"Num min. ctx tokens: {min(ctx_tokens)}")
+        print(f"Num std. ctx tokens: {np.std(ctx_tokens)}")
+        plt.hist(ctx_tokens, bins=100)
+        plt.savefig(f"tmp/ctx_n_token_hist_{split}_{ds_name}.png")
+        plt.close()

tmp/visualize_long_samples.py

@@ -0,0 +1,133 @@
+import os
+from functools import partial
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from ctx_to_lora.data.processing import get_tokenized_dataset, load_and_process_dataset
+from ctx_to_lora.model_loading import get_tokenizer
+
+tokenizer_kwargs = {"max_length": 2**13}
+ctx_tokenizer_kwargs = {"max_length": 2**13}  # not used for now
+tokenizer = get_tokenizer("google/gemma-2-2b-it", train=True)
+ctx_tokenizer = get_tokenizer("google/gemma-2-2b-it", train=True)
+
+_get_tokenized_dataset = partial(
+    get_tokenized_dataset,
+    base_model_max_len=2**13,
+    tokenizer=tokenizer,
+    tokenizer_kwargs=tokenizer_kwargs,
+    ctx_model_max_len=2**13,
+    ctx_tokenizer=ctx_tokenizer,
+    ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
+    add_ctx_to_chat=False,
+    add_repeat_prompt=False,
+    add_negative_prompt=False,
+    use_kl_loss=False,
+    # streaming=data_args.streaming,
+)
+_load_and_process_dataset = partial(
+    load_and_process_dataset,
+    add_negative_prompt=False,
+    add_repeat_prompt=False,
+    repeat_prob=0,
+    is_pretrain=False,
+    streaming=False,
+    num_proc=8,
+)
+
+train_ds_names = [
+    # "fw_qa",
+    # "fw_qa_large",
+    # "ctx_qa",
+    # "pwc",
+    # "hotpot_qa",
+    # "squad",
+    # "drop",
+    # "narrativeqa",
+    # "quoref",
+    # "ropes",
+    # "synthetic_convqa",
+    # "booksum",
+    # "gov_report",
+    # "wikitext-2",
+    # "wikitext-103",
+    # "fw_qa_2"
+    # "fw_qa_xl",
+    # "fw_qa_3_medium",
+    "self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_0_tiny",
+    "self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_1_tiny",
+    "self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_2_tiny",
+    "self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_3_tiny",
+]
+loaded_dataset = dict()
+val_ds_names = []  # ["fw_qa_large", "ctx_qa", "pwc", "hotpot_qa", "squad"]
+for split, ds_names in zip(["train", "validation"], [train_ds_names, val_ds_names]):
+    loaded_dataset[split] = {
+        os.path.basename(ds_name): _load_and_process_dataset(ds_name, split)
+        for ds_name in ds_names
+    }
+    for ds_name in loaded_dataset[split]:
+        print(ds_name)
+        print(f"Num samples: {len(loaded_dataset[split][ds_name])}")
+        ds = loaded_dataset[split][ds_name]
+        sample = loaded_dataset[split][ds_name][0]
+
+        # TODO: let's investigate samples that are very long and very short
+        if "responses" in sample:
+            char_lens = []
+            responses = []
+            questions = []
+            ctxs = []
+            for sample in loaded_dataset[split][ds_name]:
+                char_lens += [len(res) for res in sample["responses"]]
+                responses += sample["responses"]
+                questions += sample["prompts"]
+                ctxs += [sample["context"]] * len(sample["responses"])
+        else:
+            responses = ds["response"]
+            questions = ds["prompt"]
+            ctxs = ds["context"]
+
+            char_lens = [
+                len(sample["response"]) for sample in loaded_dataset[split][ds_name]
+            ]
+
+        print(f"Num avg. char_lens: {sum(char_lens) / len(char_lens)}")
+        print(f"Num max. char_lens: {max(char_lens)}")
+        print(f"Num min. char_lens: {min(char_lens)}")
+        print(f"Num std. char_lens: {np.std(char_lens)}")
+        plt.hist(char_lens, bins=100)
+        plt.savefig(f"tmp/n_token_hist_{split}_{ds_name}.png")
+        plt.close()
+
+        long_indices = [i for i, clen in enumerate(char_lens) if clen > 2000]
+        short_indices = [i for i, clen in enumerate(char_lens) if clen < 100]
+        breakpoint()
+        for i in long_indices:
+            print(f"Long sample {i}")
+            print(f"Response length: {len(responses[i])}")
+            print(f"Context: {ctxs[i]}")
+            print(f"Question: {questions[i]}")
+            print(f"Response: {responses[i]}")
+            print("=" * 80)
+
+        breakpoint()
+        for i in short_indices:
+            print(f"Short sample {i}")
+            print(f"Response length: {len(responses[i])}")
+            print(f"Context: {ctxs[i]}")
+            print(f"Question: {questions[i]}")
+            print(f"Response: {responses[i]}")
+            print("=" * 80)
+
+        ctx_tokens = [
+            len(sample["ctx_ids"]) for sample in loaded_dataset[split][ds_name]
+        ]
+        print(f"Num avg. ctx char_lens: {sum(ctx_tokens) / len(ctx_tokens)}")
+        print(f"Num max. ctx char_lens: {max(ctx_tokens)}")
+        print(f"Num min. ctx char_lens: {min(ctx_tokens)}")
+        print(f"Num std. ctx char_lens: {np.std(ctx_tokens)}")
+        plt.hist(ctx_tokens, bins=100)
+        plt.savefig(f"tmp/ctx_n_token_hist_{split}_{ds_name}.png")
+        plt.close()