Finalize llama3 example

README.md

@@ -27,9 +27,6 @@ For detailed description of the methodology, see the [**paper**](https://arxiv.o
 
 ---
 
-> [!WARNING]
-> This repository is under development and has not reached its first stable release.
-
 ## Installation
 
 > [!IMPORTANT]

examples/openwebtext/compute_scores.py

@@ -1,12 +1,9 @@
 import argparse
 import logging
 from datetime import timedelta
-from typing import Dict, List, Optional
 
 import torch
-import torch.nn.functional as F
 from accelerate import Accelerator, InitProcessGroupKwargs
-from torch import nn
 from transformers import default_data_collator
 
 from examples.openwebtext.pipeline import (
@@ -16,18 +13,11 @@
 )
 from examples.openwebtext.task import LanguageModelingTask
 from kronfluence.analyzer import Analyzer, prepare_model
-from kronfluence.task import Task
-from kronfluence.utils.common.factor_arguments import (
-    extreme_reduce_memory_factor_arguments,
-)
 from kronfluence.utils.common.score_arguments import (
-    all_low_precision_score_arguments,
     extreme_reduce_memory_score_arguments,
 )
 from kronfluence.utils.dataset import DataLoaderKwargs
 
-BATCH_TYPE = Dict[str, torch.Tensor]
-
 torch.backends.cudnn.benchmark = True
 torch.backends.cuda.matmul.allow_tf32 = True
 
@@ -36,9 +26,9 @@ def parse_args():
     parser = argparse.ArgumentParser(description="Influence score computation on Openwebtext dataset.")
 
     parser.add_argument(
-        "--factor_strategy",
+        "--factors_name",
         type=str,
-        default="ekfac",
+        default="july_11",
         help="Strategy to compute influence factors.",
     )
     parser.add_argument(
@@ -50,7 +40,7 @@ def parse_args():
     parser.add_argument(
         "--train_batch_size",
         type=int,
-        default=4,
+        default=8,
         help="Batch size for computing query gradients.",
     )
     parser.add_argument(
@@ -93,22 +83,24 @@ def main():
     dataloader_kwargs = DataLoaderKwargs(num_workers=4, collate_fn=default_data_collator, pin_memory=True)
     analyzer.set_dataloader_kwargs(dataloader_kwargs)
 
-    scores_name = args.factor_strategy
+    scores_name = args.factors_name
     rank = args.query_gradient_rank if args.query_gradient_rank != -1 else None
+    # We set the damping term used for LLMs.
     score_args = extreme_reduce_memory_score_arguments(
         damping_factor=None, module_partitions=1, query_gradient_low_rank=rank, dtype=torch.bfloat16
     )
-    # score_args.module_partitions = 2
     score_args.query_gradient_accumulation_steps = 10
+    # We can invest some time in getting more accurate SVD results.
+    score_args.use_full_svd = True
     analyzer.compute_pairwise_scores(
         scores_name=scores_name,
         score_args=score_args,
-        factors_name=args.factor_strategy,
+        factors_name=args.factors_name,
         query_dataset=eval_dataset,
         train_dataset=train_dataset,
         per_device_query_batch_size=1,
         per_device_train_batch_size=args.train_batch_size,
-        overwrite_output_dir=False,
+        overwrite_output_dir=True,
     )
     scores = analyzer.load_pairwise_scores(scores_name)["all_modules"]
     logging.info(f"Scores shape: {scores.shape}")

examples/openwebtext/files/factor_arguments.json

@@ -0,0 +1,20 @@
+{
+    "strategy": "ekfac",
+    "use_empirical_fisher": false,
+    "amp_dtype": "torch.bfloat16",
+    "amp_scale": 65536.0,
+    "has_shared_parameters": false,
+    "covariance_max_examples": 100000,
+    "covariance_data_partitions": 1,
+    "covariance_module_partitions": 2,
+    "activation_covariance_dtype": "torch.bfloat16",
+    "gradient_covariance_dtype": "torch.bfloat16",
+    "eigendecomposition_dtype": "torch.float64",
+    "lambda_max_examples": 100000,
+    "lambda_data_partitions": 1,
+    "lambda_module_partitions": 4,
+    "use_iterative_lambda_aggregation": true,
+    "offload_activations_to_cpu": true,
+    "per_sample_gradient_dtype": "torch.bfloat16",
+    "lambda_dtype": "torch.bfloat16"
+}

examples/openwebtext/files/query_dataset_metadata.json

@@ -0,0 +1,5 @@
+{
+    "type": "Dataset",
+    "dataset_size": 5,
+    "indices": null
+}

examples/openwebtext/files/score_arguments.json

@@ -0,0 +1,19 @@
+{
+    "damping_factor": null,
+    "amp_dtype": "torch.bfloat16",
+    "offload_activations_to_cpu": true,
+    "data_partitions": 1,
+    "module_partitions": 1,
+    "compute_per_module_scores": false,
+    "compute_per_token_scores": false,
+    "query_gradient_accumulation_steps": 10,
+    "query_gradient_low_rank": 64,
+    "use_full_svd": true,
+    "aggregate_query_gradients": false,
+    "aggregate_train_gradients": false,
+    "use_measurement_for_self_influence": false,
+    "query_gradient_svd_dtype": "torch.float32",
+    "per_sample_gradient_dtype": "torch.bfloat16",
+    "precondition_dtype": "torch.bfloat16",
+    "score_dtype": "torch.bfloat16"
+}

examples/openwebtext/files/train_dataset_metadata.json

@@ -0,0 +1,5 @@
+{
+    "type": "Dataset",
+    "dataset_size": 100000,
+    "indices": null
+}

examples/openwebtext/fit_factors.py

@@ -1,12 +1,9 @@
 import argparse
 import logging
 from datetime import timedelta
-from typing import Dict, List, Optional
 
 import torch
-import torch.nn.functional as F
 from accelerate import Accelerator, InitProcessGroupKwargs
-from torch import nn
 from transformers import default_data_collator
 
 from examples.openwebtext.pipeline import construct_llama3, get_openwebtext_dataset
@@ -17,8 +14,6 @@
 )
 from kronfluence.utils.dataset import DataLoaderKwargs
 
-BATCH_TYPE = Dict[str, torch.Tensor]
-
 torch.backends.cudnn.benchmark = True
 torch.backends.cuda.matmul.allow_tf32 = True
 

examples/openwebtext/inpsect_factors.py

@@ -0,0 +1,39 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.colors import LogNorm
+from tueplots import markers
+
+from kronfluence.analyzer import Analyzer
+
+
+def main():
+    plt.rcParams.update({"figure.dpi": 150})
+    plt.rcParams.update(markers.with_edge())
+    plt.rcParams["axes.axisbelow"] = True
+
+    layer_num = 18
+    module_name = f"model.layers.{layer_num}.mlp.down_proj"
+    # module_name = f"model.layers.{layer_num}.mlp.up_proj"
+    lambda_processed = Analyzer.load_file("num_lambda_processed.safetensors")[module_name]
+    lambda_matrix = Analyzer.load_file("lambda_matrix.safetensors")[module_name]
+    lambda_matrix.div_(lambda_processed)
+    lambda_matrix = lambda_matrix.float()
+    plt.matshow(lambda_matrix, cmap="PuBu", norm=LogNorm())
+
+    plt.title(module_name)
+    plt.colorbar()
+    plt.show()
+    plt.clf()
+
+    lambda_matrix = lambda_matrix.view(-1).numpy()
+    sorted_lambda_matrix = np.sort(lambda_matrix)
+    plt.plot(sorted_lambda_matrix)
+    plt.title(module_name)
+    plt.grid()
+    plt.yscale("log")
+    plt.ylabel("Eigenvalues")
+    plt.show()
+
+
+if __name__ == "__main__":
+    main()

examples/openwebtext/inspect_scores.py

@@ -0,0 +1,42 @@
+import matplotlib.pyplot as plt
+import torch
+from transformers import AutoTokenizer
+
+from examples.openwebtext.pipeline import (
+    MODEL_NAME,
+    get_custom_dataset,
+    get_openwebtext_dataset,
+)
+from kronfluence.analyzer import Analyzer
+
+
+def main():
+    scores = Analyzer.load_file("influence_results/scores_jul_11_2024/pairwise_scores.safetensors")["all_modules"].float()
+
+    train_dataset = get_openwebtext_dataset()
+    eval_dataset = get_custom_dataset()
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True, trust_remote_code=True)
+
+    eval_idx = 4
+    sorted_scores = torch.sort(scores[eval_idx], descending=True)
+    top_indices = sorted_scores.indices
+
+    plt.plot(sorted_scores.values)
+    plt.grid()
+    plt.ylabel("IF Score")
+    plt.show()
+
+    print("Query Sequence:")
+    print(
+        "Prompt: " + eval_dataset[eval_idx]["prompt"] + "; Completion: " + eval_dataset[eval_idx]["completion"] + "\n"
+    )
+
+    print("Top Influential Sequences:")
+    for i in range(100):
+        print("=" * 80)
+        print(f"Rank = {i}; Score = {scores[eval_idx][int(top_indices[i])].item()}")
+        print(tokenizer.decode(train_dataset[int(top_indices[i])]["input_ids"]))
+
+
+if __name__ == "__main__":
+    main()

examples/openwebtext/requirements.txt

@@ -1,2 +1,4 @@
 transformers
 datasets
+matplotlib
+tueplots

kronfluence/arguments.py

@@ -151,6 +151,10 @@ def __post_init__(self) -> None:
         ):
             raise ValueError("All data and module partitions must be positive.")
 
+        # For backward compatibility:
+        if not hasattr(self, "amp_scale"):
+            self.amp_scale = 2.0**16
+
 
 @dataclass
 class ScoreArguments(Arguments):

kronfluence/factor/config.py

@@ -9,6 +9,7 @@
     ACTIVATION_EIGENVECTORS_NAME,
     GRADIENT_EIGENVALUES_NAME,
     GRADIENT_EIGENVECTORS_NAME,
+    HEURISTIC_DAMPING_SCALE,
     LAMBDA_MATRIX_NAME,
     NUM_LAMBDA_PROCESSED,
 )
@@ -199,7 +200,7 @@ def prepare(self, storage: STORAGE_TYPE, score_args: Any, device: torch.device)
         lambda_matrix.div_(storage[NUM_LAMBDA_PROCESSED].to(device=device))
         damping_factor = score_args.damping_factor
         if damping_factor is None:
-            damping_factor = 0.1 * torch.mean(lambda_matrix)
+            damping_factor = HEURISTIC_DAMPING_SCALE * torch.mean(lambda_matrix)
         lambda_matrix.add_(damping_factor)
         storage[LAMBDA_MATRIX_NAME] = lambda_matrix.to(dtype=score_args.precondition_dtype, device="cpu").contiguous()
         storage[NUM_LAMBDA_PROCESSED] = None
@@ -259,7 +260,7 @@ def prepare(self, storage: STORAGE_TYPE, score_args: Any, device: torch.device)
         lambda_matrix = torch.kron(activation_eigenvalues.unsqueeze(0), gradient_eigenvalues.unsqueeze(-1)).unsqueeze(0)
         damping_factor = score_args.damping_factor
         if damping_factor is None:
-            damping_factor = 0.1 * torch.mean(lambda_matrix)
+            damping_factor = HEURISTIC_DAMPING_SCALE * torch.mean(lambda_matrix)
         lambda_matrix.add_(damping_factor)
         storage[LAMBDA_MATRIX_NAME] = lambda_matrix.to(dtype=score_args.precondition_dtype, device="cpu").contiguous()
         storage[NUM_LAMBDA_PROCESSED] = None
@@ -328,7 +329,7 @@ def prepare(self, storage: STORAGE_TYPE, score_args: Any, device: torch.device)
         lambda_matrix.div_(storage[NUM_LAMBDA_PROCESSED].to(device=device))
         damping_factor = score_args.damping_factor
         if damping_factor is None:
-            damping_factor = 0.1 * torch.mean(lambda_matrix)
+            damping_factor = HEURISTIC_DAMPING_SCALE * torch.mean(lambda_matrix)
         lambda_matrix.add_(damping_factor)
         storage[LAMBDA_MATRIX_NAME] = lambda_matrix.to(dtype=score_args.precondition_dtype, device="cpu").contiguous()
         storage[NUM_LAMBDA_PROCESSED] = None

kronfluence/utils/constants.py

@@ -18,6 +18,8 @@
 # The total iteration step to synchronize the process when using distributed setting.
 DISTRIBUTED_SYNC_INTERVAL = 1_000
 
+HEURISTIC_DAMPING_SCALE = 0.1
+
 # Activation covariance matrix.
 ACTIVATION_COVARIANCE_MATRIX_NAME = "activation_covariance"
 # Pseudo-gradient covariance matrix.

tests/test_analyzer.py

@@ -104,6 +104,7 @@ def test_default_factor_arguments() -> None:
     assert factor_args.strategy == "ekfac"
     assert factor_args.use_empirical_fisher is False
     assert factor_args.amp_dtype is None
+    assert factor_args.amp_scale == 2.0**16
     assert factor_args.has_shared_parameters is False
 
     assert factor_args.covariance_max_examples == 100_000