src/transformer_analysis/weight_analysis.py

# `run_weight_analysis.py`
# Analysis main for weight analysis
import json
import logging
import uuid
import os
import shutil
from datetime import datetime
from types import SimpleNamespace

from tqdm import tqdm
import pandas as pd
from huggingface_hub import snapshot_download
from datasets import Dataset
from transformers import AutoConfig

from transformer_analysis.perf_logger import PerfLogger
from transformer_analysis.attn_head_analysis import LayerHeadContainer
from transformer_analysis.histogram_utils import (
    stats_config_default,
    weight_bins_default,
    sv_bins_default,
    make_weight_bins,
    make_sv_bins,
)
from transformer_analysis.model_registry import (
    get_model_config,
    extract_weight_map,
)
from transformer_analysis.device_utils import get_device


def process_model(
    model_name="pythia-70m-deduped",
    revision=None,
    idx_max=-1,
    out_dir="histos",
    cache_dir="./model_data",
    cleanup_downloads=False,
    low_rank_svd_approximation=False,
    top_k_svd=-1,
    resume_download=True,
    max_workers=4,
    device=None,
    skip_postprocess=False,
    binning_strategy="fixed",
):
    job_uuid = str(uuid.uuid4())[:8]
    job_id = datetime.now().strftime("%Y%m%d_%H%M%S")

    # Configure logging
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(message)s",
        handlers=[
            logging.FileHandler(f"{out_dir}/logs/{job_id}.log"),
            logging.StreamHandler(),
        ],
    )

    perf = PerfLogger(job_id)

    logging.info(f"Starting job {job_id} {job_uuid}")

    device_str = str(get_device(device))
    logging.info(f"Using device: {device_str}")

    with perf.phase("load_model"):
        logging.info("Loading model...")
        model_config = get_model_config(model_name)

        revision_string = revision if revision else "main"

        cache_path = snapshot_download(
            repo_id=model_config.repo_id,
            revision=revision,
            cache_dir=f"{cache_dir}/{model_name}/{revision_string}",
            allow_patterns=model_config.allow_patterns,
            resume_download=resume_download,
            max_workers=max_workers,
        )
        hf_config = AutoConfig.from_pretrained(cache_path)

    logging.info(perf.log_report(context=model_name))

    # Phase 2: Configuration
    with perf.phase("configure"):
        logging.info("Configuring analysis...")

        config = SimpleNamespace()
        config.weight_type = ["W_Q", "W_K", "W_QK", "W_Q_gram", "W_K_gram", "QK_alignment"]
        config.stats = stats_config_default.copy()
        config.w_bins = make_weight_bins(strategy=binning_strategy)
        config.sv_bins = make_sv_bins(strategy=binning_strategy)
        config.binning_strategy = binning_strategy
        config.use_density = True
        config.n_heads = model_config.get_config_value(hf_config.__dict__, "n_heads")
        config.d_model = model_config.get_config_value(hf_config.__dict__, "d_model")
        config.n_layers = model_config.get_config_value(hf_config.__dict__, "n_layers")
        config.head_dim = config.d_model // config.n_heads
        # SVD configuration options (passed from function parameters)
        config.low_rank_svd_approximation = low_rank_svd_approximation
        config.top_k_svd = top_k_svd

        n_layers, n_heads, head_dim = config.n_layers, config.n_heads, config.head_dim
        d_model = config.d_model
    logging.info(perf.log_report())

    # Phase 3: Loop with conditional logging
    layer_data = []
    with perf.phase("loop"):
        n_hl = n_heads * n_layers
        if idx_max == -1:
            idx_max = n_hl
        else:
            idx_max = min(abs(idx_max), n_hl)
        logging.info(f"Processing {idx_max} / {n_hl}")

        # Get weight_map, needed if safetensors format unavailable and bin files are sharded
        weight_map = extract_weight_map(cache_path=cache_path)

        for layer_idx in tqdm(range(n_layers), desc="Processing layers", leave=True):
            W_Q, W_K, _ = model_config.extract_qkv(
                cache_path, layer_idx, d_model, weight_map, device=device_str, qkv_scale_factor=model_config.qkv_scale_factor
            )

            # For per-head analysis:
            W_Q_h = W_Q.reshape(n_heads, head_dim, d_model).float()
            W_K_h = W_K.reshape(n_heads, head_dim, d_model).float()

            hc = LayerHeadContainer(
                layer_idx,
                config,
                low_rank_svd_approximation=config.low_rank_svd_approximation,
                top_k_svd=config.top_k_svd,
                device=device_str
            )
            layer_input = {"W_Q": W_Q_h, "W_K": W_K_h}
            hc.analyze_layer(layer_input)
            layer_data.append(hc)
            del W_Q, W_K, W_Q_h, W_K_h
    logging.info(perf.log_report())

    # Phase 4: Finalization
    dfs = []
    with perf.phase("finalize"):
        logging.info("Aggregating results...")
        for lhc in layer_data:
            if not skip_postprocess:
                lhc.post_process()
            dfs.append(lhc.to_pandas())
        df = pd.concat(dfs, ignore_index=True)
        df["model"] = model_name
        if revision:
            df["revision"] = revision
            df["step"] = int(revision.strip("step"))
        df["job_uuid"] = job_uuid
        df["job_id"] = job_id
    logging.info(perf.log_report())

    # Phase 5: Write output
    with perf.phase("write_output"):
        logging.info("Writing outputs...")

        ds = Dataset.from_pandas(df)
        ds.info.description = "metadata.json"
        out_prefix = f"{out_dir}/{model_name}_{revision_string}"
        ds.save_to_disk(out_prefix)
        logging.info(f"Saving dataset: {out_prefix}")
        # for metadata we need to do some coversions to make the objects JSON serializable
        config_dict = vars(config).copy()
        config_dict["stats"] = {k: v.__name__ for k, v in config_dict["stats"].items()}
        config_dict["w_bins"] = config_dict["w_bins"].tolist()
        config_dict["sv_bins"] = config_dict["sv_bins"].tolist()

        with open(f"{out_prefix}/metadata.json", "w") as f:
            json.dump(config_dict, f, indent=2)
        with open(f"{out_dir}/logs/perf_{job_id}.json", "w") as f:
            json.dump(perf.to_metadata(), f, indent=2)
    logging.info(perf.log_report())

    # Phase 6: Cleanup
    with perf.phase("cleanup"):
        if cleanup_downloads:
            logging.info("Cleaning up downloads...")
            cache_path = f"{cache_dir}/{model_name}/{revision}"
            if os.path.exists(cache_path):
                shutil.rmtree(cache_path)

    disk_usage = shutil.disk_usage(cache_dir)
    logging.info(
        f"Disk usage: {disk_usage.used / (1024**3):.2f} GB / {disk_usage.total / (1024**3):.2f} GB"
    )
    logging.info(perf.log_report())

    # Summary
    logging.info("\n" + "=" * 60)
    logging.info("Performance Summary:")
    for phase_name in perf.phases.keys():
        logging.info(perf.log_report(phase=phase_name))
    logging.info("=" * 60)
    logging.info(f"Performance saved to {out_dir}/logs/perf_{job_id}.json")


def reprocess_metrics(
    model_name: str,
    revision=None,
    all_revisions: bool = False,
    out_dir: str = "outputs",
    quiet: bool = False,
):
    """
    Reprocess existing datasets to update/add metrics without re-running full analysis.

    This function loads existing datasets, recomputes metrics (both weight histogram
    and singular value metrics), and overwrites the dataset with updated columns.

    Args:
        model_name: Name of the model to reprocess
        revision: Specific revision to reprocess (or None for main)
        all_revisions: Whether to process all available revisions
        out_dir: Output directory containing existing datasets
        quiet: Whether to suppress output
    """
    from datasets import load_from_disk, Dataset
    from transformer_analysis.histogram_utils import (
        normality_metrics,
        singular_value_metrics,
        get_model_versions,
    )
    import numpy as np

    if not quiet:
        print("\n" + "=" * 80)
        print(f"Reprocessing Metrics: {model_name}")
        print("=" * 80)

    # Determine which revisions to process
    if all_revisions:
        revisions = get_model_versions(model_name)
        if not revisions:
            print(f"Model {model_name} has no revisions defined. Processing main branch only.")
            revisions = [None]
    elif revision:
        revisions = [revision]
    else:
        revisions = [None]

    if not quiet:
        print(f"Revisions to process: {len(revisions)}")

    for rev in tqdm(revisions, desc=f"Reprocessing {model_name}", disable=quiet):
        revision_str = rev if rev else "main"

        # Determine dataset path
        if rev:
            dataset_path = os.path.join(out_dir, f"{model_name}_{revision_str}")
        else:
            dataset_path = os.path.join(out_dir, model_name)

        if not os.path.exists(dataset_path):
            print(f"  WARNING: Dataset not found at {dataset_path}, skipping...")
            continue

        if not quiet:
            print(f"\n  Reprocessing: {model_name} @ {revision_str}")

        try:
            # Load existing dataset
            ds = load_from_disk(dataset_path)
            df = ds.to_pandas()

            # Load metadata to get bin information
            metadata_path = os.path.join(dataset_path, "metadata.json")
            with open(metadata_path, "r") as f:
                metadata = json.load(f)

            w_bins = np.array(metadata["w_bins"])
            centers = (w_bins[:-1] + w_bins[1:]) / 2

            if not quiet:
                print(f"    Loaded dataset with {len(df)} rows")

            # Process each row
            new_columns = {}
            for idx, row in tqdm(df.iterrows(), total=len(df), desc="  Processing rows", disable=quiet, leave=False):
                # Create a dictionary for this row (simulating the 'h' dict)
                h = row.to_dict()

                # Apply weight histogram metrics
                for metric_func in normality_metrics.values():
                    metric_func(h, centers)

                # Apply singular value metrics if SVD data exists
                if "SVD" in h and h["SVD"] is not None and not pd.isna(h["SVD"]).all():
                    svd_array = h["SVD"]
                    if isinstance(svd_array, (list, np.ndarray)) and len(svd_array) > 0:
                        for metric_func in singular_value_metrics.values():
                            metric_func(h, svd_array)

                # Store new metric values
                for key, value in h.items():
                    if key not in row or row[key] != value:
                        if key not in new_columns:
                            new_columns[key] = [None] * len(df)
                        new_columns[key][idx] = value

            # Add new columns to dataframe
            for col_name, col_values in new_columns.items():
                df[col_name] = col_values
                if not quiet:
                    print(f"    Added/updated column: {col_name}")

            # Save updated dataset
            updated_ds = Dataset.from_pandas(df)
            updated_ds.info.description = "metadata.json"
            updated_ds.save_to_disk(dataset_path)

            if not quiet:
                print(f"    ✓ Saved updated dataset to {dataset_path}")

        except Exception as e:
            print(f"  ERROR reprocessing {model_name} @ {revision_str}: {e}")
            import traceback
            traceback.print_exc()
            continue

    if not quiet:
        print("\n" + "=" * 80)
        print(f"Completed reprocessing: {model_name}")
        print("=" * 80 + "\n")


def create_campaign(path, name, clobber=False, logs=True):
    base_dir = os.path.join(path, name)
    log_dir = os.path.join(base_dir, "logs")

    if not os.path.exists(base_dir):
        os.makedirs(base_dir)
        os.makedirs(log_dir)
        return base_dir

    elif clobber:
        shutil.rmtree(base_dir)
        os.makedirs(base_dir)
        os.makedirs(log_dir)
        return base_dir

    return base_dir


def write_dataset_and_metadata(ds_list, metadata, ds_name):
    """Write combined dataset and metadata to disk.

    Args:
        ds_list: List of datasets to concatenate
        metadata: Metadata dictionary to write
        ds_name: Output directory name for the dataset
    """
    from datasets import concatenate_datasets

    combined_ds = concatenate_datasets(ds_list)
    combined_ds.info.description = "metadata.json"
    combined_ds.save_to_disk(ds_name)
    with open(f"{ds_name}/metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)


def merge_versions(
    model_name="pythia-70m-deduped", path="histos", suffix="all_checkpoints"
):
    """Merge all checkpoint versions of a single model into one dataset.

    Args:
        model_name: Name of the model to merge
        path: Base directory containing the datasets
        suffix: Suffix for the output merged dataset
    """
    from datasets import load_from_disk
    from transformer_analysis.histogram_utils import get_model_versions

    ds_list = []
    metadata = None
    for rev in tqdm(get_model_versions(model_name), desc=f"Processing {model_name}"):
        pattern = f"{model_name}_{rev}"
        ds = load_from_disk(f"{path}/{pattern}")
        ds_list.append(ds)
        if metadata is None:
            with open(f"{path}/{pattern}/{ds.info.description}") as f:
                metadata = json.load(f)
    write_dataset_and_metadata(ds_list, metadata, f"{path}/{model_name}_{suffix}")


def merge_datasets(model_name_list, path="histos", out_name="merged", suffix=None):
    """Merge multiple model datasets into a single combined dataset.

    Args:
        model_name_list: List of model names (or patterns) to merge
        path: Base directory containing the datasets
        out_name: Name for the output merged dataset
        suffix: Optional suffix to append to each model name pattern
    """
    from datasets import load_from_disk

    META_MERGE_KEY = "merged"

    ds_list = []
    combined_metadata = None
    merged_dict = {}
    for model_name in tqdm(model_name_list, desc="Processing models"):
        pattern = model_name
        if suffix is not None and isinstance(str, suffix):
            pattern += "_" + suffix
        ds = load_from_disk(f"{path}/{pattern}")
        ds_list.append(ds)

        # now the metadata
        mf = f"{path}/{pattern}/{ds.info.description}"
        with open(mf) as f:
            metadata = json.load(f)
        if combined_metadata is None:
            combined_metadata = {
                k: v for k, v in metadata.items() if k != META_MERGE_KEY
            }
        model_name = model_name.rstrip("_main")
        if META_MERGE_KEY in metadata:  # Flatten
            for k, v in metadata[META_MERGE_KEY].items():
                key = k
                while key in merged_dict:  # make key name unique
                    key = f"{model_name}_{key}"
                merged_dict[key] = v
        else:
            merged_dict[model_name] = metadata

    combined_metadata.update({META_MERGE_KEY: merged_dict})
    write_dataset_and_metadata(ds_list, combined_metadata, f"{path}/{out_name}")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default="gpt2")  # "pythia-70m-deduped")
    parser.add_argument("--out", type=str, default="Drive/ana-002")
    parser.add_argument("--cache", type=str, default="./model_data")
    parser.add_argument("--clobber", type=bool, default=False)
    parser.add_argument("--test", action="store_true", default=False)
    parser.add_argument("--low-rank-svd", action="store_true", default=False, dest="low_rank_svd")
    parser.add_argument("--top-k-svd", type=int, default=-1, dest="top_k_svd")
    parser.add_argument("--resume-download", action="store_true", default=True, dest="resume_download")
    parser.add_argument("--no-resume-download", action="store_false", dest="resume_download")
    parser.add_argument("--max-workers", type=int, default=4, dest="max_workers")
    parser.add_argument("--device", type=str, default=None, choices=["cuda", "mps", "cpu"])

    args = parser.parse_args()
    if args.test:
        print("=" * 20 + "Test option selected" + "=" * 20)
        print("\t\t" + "output and clobber options will be overwritten")
        args.out, args.clobber = "test", True
    cwd = os.getcwd()
    out_dir = create_campaign(path=cwd, name=args.out, clobber=args.clobber, logs=True)
    model_name = args.model

    model_config = get_model_config(args.model)
    revisions = model_config.revisions
    if args.test:
        revisions = revisions[-1:] if revisions else None
    else:
        from transformers import logging as hf_logging

        hf_logging.set_verbosity_error()
        import warnings

        warnings.filterwarnings("ignore")

    # loop on checkpoints
    if revisions:
        for revision in tqdm(revisions):
            process_model(
                model_name=model_name,
                revision=revision,
                out_dir=out_dir,
                cache_dir=args.cache,
                low_rank_svd_approximation=args.low_rank_svd,
                top_k_svd=args.top_k_svd,
                resume_download=args.resume_download,
                max_workers=args.max_workers,
                device=args.device,
            )
    else:
        process_model(
            model_name=model_name,
            revision=None,
            out_dir=out_dir,
            cache_dir=args.cache,
            low_rank_svd_approximation=args.low_rank_svd,
            top_k_svd=args.top_k_svd,
            resume_download=args.resume_download,
            max_workers=args.max_workers,
            device=args.device,
        )