bergson 0.0.1__py3-none-any.whl

bergson/__init__.py

@@ -0,0 +1,21 @@
+__version__ = "0.0.1"
+
+from .attributor import Attributor
+from .collection import collect_gradients
+from .data import DataConfig, IndexConfig, load_gradients
+from .faiss_index import FaissConfig
+from .gradcheck import FiniteDiff
+from .gradients import GradientCollector, GradientProcessor, HeadConfig
+
+__all__ = [
+    "collect_gradients",
+    "load_gradients",
+    "Attributor",
+    "FaissConfig",
+    "FiniteDiff",
+    "GradientCollector",
+    "GradientProcessor",
+    "IndexConfig",
+    "DataConfig",
+    "HeadConfig",
+]

bergson/__main__.py

@@ -0,0 +1,12 @@
+from simple_parsing import parse
+
+from .build import build_gradient_dataset
+from .data import IndexConfig
+
+
+def main():
+    build_gradient_dataset(parse(IndexConfig))
+
+
+if __name__ == "__main__":
+    main()

bergson/attributor.py

@@ -0,0 +1,159 @@
+from collections import defaultdict
+from contextlib import contextmanager
+from typing import Generator
+
+import torch
+from torch import Tensor, nn
+
+from .data import load_gradients
+from .faiss_index import FaissConfig, FaissIndex
+from .gradients import GradientCollector, GradientProcessor
+
+
+class TraceResult:
+    """Result of a .trace() call."""
+
+    def __init__(self):
+        # Should be set by the Attributor after a search
+        self._indices: Tensor | None = None
+        self._scores: Tensor | None = None
+
+    @property
+    def indices(self) -> Tensor:
+        """The indices of the top-k examples."""
+        if self._indices is None:
+            raise ValueError("No indices available. Exit the context manager first.")
+
+        return self._indices
+
+    @property
+    def scores(self) -> Tensor:
+        """The attribution scores of the top-k examples."""
+        if self._scores is None:
+            raise ValueError("No scores available. Exit the context manager first.")
+
+        return self._scores
+
+
+class Attributor:
+    def __init__(
+        self,
+        index_path: str,
+        device: str = "cpu",
+        dtype: torch.dtype = torch.float32,
+        unit_norm: bool = False,
+        faiss_cfg: FaissConfig | None = None,
+    ):
+        self.device = device
+        self.dtype = dtype
+        self.unit_norm = unit_norm
+        self.faiss_index = None
+
+        # Load the gradient processor
+        self.processor = GradientProcessor.load(index_path, map_location=device)
+
+        # Load the gradient index
+        if faiss_cfg:
+            self.faiss_index = FaissIndex(index_path, faiss_cfg, device, unit_norm)
+            self.N = self.faiss_index.ntotal
+        else:
+            mmap = load_gradients(index_path)
+
+            # Copy gradients into device memory
+            self.grads = {
+                name: torch.tensor(mmap[name], device=device, dtype=dtype)
+                for name in mmap.dtype.names
+            }
+            self.N = mmap[mmap.dtype.names[0]].shape[0]
+
+            if unit_norm:
+                norm = torch.cat([grad for grad in self.grads.values()], dim=1).norm(
+                    dim=1, keepdim=True
+                )
+                for name in self.grads:
+                    self.grads[name] /= norm
+
+    def search(
+        self, queries: dict[str, Tensor], k: int, modules: list[str] | None = None
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Search for the `k` nearest examples in the index based on the query or queries.
+
+        Args:
+            queries: The query tensor of shape [..., d].
+            k: The number of nearest examples to return for each query.
+            module: The name of the module to search for. If `None`,
+                all modules will be searched.
+
+        Returns:
+            A namedtuple containing the top `k` indices and inner products for each
+            query. Both have shape [..., k].
+        """
+        q = {name: item.to(self.device, self.dtype) for name, item in queries.items()}
+
+        if self.unit_norm:
+            norm = torch.cat(list(q.values()), dim=1).norm(dim=1, keepdim=True)
+
+            for name in q:
+                q[name] /= norm + 1e-8
+
+        if self.faiss_index:
+            if modules:
+                raise NotImplementedError(
+                    "FAISS index does not implement module-specific search."
+                )
+
+            q = torch.cat([q[name] for name in q], dim=1).cpu().numpy()
+
+            distances, indices = self.faiss_index.search(q, k)
+
+            return torch.from_numpy(distances.squeeze()), torch.from_numpy(
+                indices.squeeze()
+            )
+
+        modules = modules or list(q.keys())
+        k = min(k, self.N)
+
+        scores = torch.stack(
+            [q[name] @ self.grads[name].mT for name in modules], dim=-1
+        ).sum(-1)
+
+        return torch.topk(scores, k)
+
+    @contextmanager
+    def trace(
+        self, module: nn.Module, k: int, *, precondition: bool = False
+    ) -> Generator[TraceResult, None, None]:
+        """
+        Context manager to trace the gradients of a module and return the
+        corresponding Attributor instance.
+        """
+        mod_grads = defaultdict(list)
+        result = TraceResult()
+
+        def callback(name: str, g: Tensor):
+            # Precondition the gradient using Cholesky solve
+            if precondition:
+                eigval, eigvec = self.processor.preconditioners_eigen[name]
+                eigval_inverse_sqrt = 1.0 / (eigval).sqrt()
+                P = eigvec * eigval_inverse_sqrt @ eigvec.mT
+                g = g.flatten(1).type_as(P)
+                g = g @ P
+            else:
+                g = g.flatten(1)
+
+            # Store the gradient for later use
+            mod_grads[name].append(g.to(self.device, self.dtype, non_blocking=True))
+
+        with GradientCollector(module, callback, self.processor):
+            yield result
+
+        if not mod_grads:
+            raise ValueError("No grads collected. Did you forget to call backward?")
+
+        queries = {name: torch.cat(g, dim=1) for name, g in mod_grads.items()}
+
+        if any(q.isnan().any() for q in queries.values()):
+            raise ValueError("NaN found in queries.")
+
+        result._scores, result._indices = self.search(queries, k)

bergson/build.py

@@ -0,0 +1,261 @@
+import os
+import socket
+from datetime import timedelta
+from typing import cast
+
+import pandas as pd
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from datasets import Dataset, IterableDataset
+from peft import PeftConfig, PeftModel
+from torch.distributed.elastic.multiprocessing import DefaultLogsSpecs, start_processes
+from torch.distributed.fsdp import fully_shard
+from tqdm.auto import tqdm
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    PreTrainedModel,
+)
+
+from .collection import collect_gradients
+from .data import IndexConfig, allocate_batches, load_data_string, tokenize
+from .gradients import GradientProcessor
+from .peft import detect_peft_modules
+from .utils import assert_type, get_layer_list
+
+
+def worker(rank: int, world_size: int, cfg: IndexConfig, ds: Dataset | IterableDataset):
+    torch.cuda.set_device(rank)
+
+    # These should be set by the main process
+    if world_size > 1:
+        addr = os.environ.get("MASTER_ADDR", "localhost")
+        port = os.environ.get("MASTER_PORT", "29500")
+
+        dist.init_process_group(
+            "nccl",
+            init_method=f"tcp://{addr}:{port}",
+            device_id=torch.device(f"cuda:{rank}"),
+            rank=rank,
+            timeout=timedelta(hours=1),
+            world_size=world_size,
+        )
+
+    match cfg.precision:
+        case "bf16":
+            dtype = torch.bfloat16
+        case "fp16":
+            dtype = torch.float16
+        case "fp32":
+            dtype = torch.float32
+        case "int4" | "int8":
+            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+        case "auto":
+            dtype = "auto"
+        case other:
+            raise ValueError(f"Unsupported precision: {other}")
+
+    device_map = {"": f"cuda:{rank}"} if not cfg.fsdp else "cpu"
+    quantization_config = None
+    if cfg.precision in ("int4", "int8"):
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=cfg.precision == "int4",
+            load_in_8bit=cfg.precision == "int8",
+            bnb_4bit_compute_dtype=dtype,
+            bnb_4bit_quant_storage=dtype,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True,
+        )
+
+    # Try to detect PEFT model
+    try:
+        peft_config = PeftConfig.from_pretrained(cfg.model)
+    except ValueError:
+        peft_config = None
+
+    if peft_config is None:
+        # Load regular model
+        model = AutoModelForCausalLM.from_pretrained(
+            cfg.model,
+            device_map=device_map,
+            quantization_config=quantization_config,
+            dtype=dtype,
+            revision=cfg.revision,
+        )
+        target_modules = None
+
+    else:
+        # Load PEFT model
+        base_model = AutoModelForCausalLM.from_pretrained(
+            peft_config.base_model_name_or_path,  # type: ignore
+            device_map=device_map,
+            quantization_config=quantization_config,
+            dtype=dtype,
+            revision=cfg.revision,
+        )
+
+        model = PeftModel.from_pretrained(
+            base_model,
+            cfg.model,
+            device_map=device_map,
+            autocast_adapter_dtype=False,
+        )
+        target_modules = detect_peft_modules(model)
+
+        # Hack for type checking
+        model = cast(PreTrainedModel, model)
+
+    if rank == 0:
+        print(f"Model loaded with dtype: {model.dtype}")
+
+    embed = model.get_input_embeddings()
+    model.requires_grad_(False)  # Freeze the model
+    embed.requires_grad_(True)  # Make sure backward hooks are called though
+
+    if cfg.fsdp:
+        # Shard each individual transformer layer
+        for layer in get_layer_list(model):
+            fully_shard(layer)
+
+        # Shard the entire model
+        fully_shard(model)
+
+    if os.path.exists(cfg.processor_path):
+        if rank == 0:
+            print(f"Loading processor from '{cfg.processor_path}'")
+
+        processor = GradientProcessor.load(
+            cfg.processor_path,
+            map_location=f"cuda:{rank}",
+        )
+    else:
+        processor = GradientProcessor(
+            {},
+            projection_dim=cfg.projection_dim or None,
+            reshape_to_square=cfg.reshape_to_square,
+            projection_type=cfg.projection_type,
+        )
+        if rank == 0:
+            processor.save(cfg.run_path)
+
+    if isinstance(ds, Dataset):
+        batches = allocate_batches(ds["length"][:], cfg.token_batch_size)
+        collect_gradients(
+            model,
+            ds,
+            processor,
+            cfg.run_path,
+            batches=batches,
+            kl_divergence=cfg.loss_fn == "kl",
+            loss_reduction=cfg.loss_reduction,
+            skip_preconditioners=cfg.skip_preconditioners,
+            target_modules=target_modules,
+            head_cfgs=cfg.head_cfgs,
+        )
+    else:
+        # Convert each shard to a Dataset then collect its gradients
+        buf, shard_id = [], 0
+
+        def flush():
+            nonlocal buf, shard_id
+            if not buf:
+                return
+            ds_shard = assert_type(Dataset, Dataset.from_list(buf))
+            batches = allocate_batches(ds_shard["length"][:], cfg.token_batch_size)
+            collect_gradients(
+                model,
+                ds_shard,
+                processor,
+                os.path.join(cfg.run_path, f"shard-{shard_id:05d}"),
+                batches=batches,
+                kl_divergence=cfg.loss_fn == "kl",
+                loss_reduction=cfg.loss_reduction,
+                skip_preconditioners=cfg.skip_preconditioners,
+                target_modules=target_modules,
+                head_cfgs=cfg.head_cfgs,
+            )
+            buf.clear()
+            shard_id += 1
+
+        for ex in tqdm(ds, desc="Collecting gradients"):
+            buf.append(ex)
+            if len(buf) == cfg.stream_shard_size:
+                flush()
+        flush()
+
+
+def dist_worker(rank: int, world_size: int, cfg: IndexConfig, ds: Dataset):
+    try:
+        worker(rank, world_size, cfg, ds)
+    finally:
+        dist.destroy_process_group()
+
+
+def estimate_advantage(ds: Dataset, cfg: IndexConfig):
+    """Group rollouts by prompt and estimate advantages."""
+    assert isinstance(ds, Dataset), "Dataset required for advantage estimation"
+
+    df = ds.select_columns([cfg.data.prompt_column, cfg.data.reward_column]).to_pandas()
+    df = assert_type(pd.DataFrame, df)
+
+    advantages = df[cfg.data.reward_column] - df.groupby(cfg.data.prompt_column)[
+        cfg.data.reward_column
+    ].transform("mean")
+
+    return advantages.tolist()
+
+
+def build_gradient_dataset(cfg: IndexConfig):
+    # In many cases the token_batch_size may be smaller than the max length allowed by
+    # the model. If cfg.data.truncation is True, we use the tokenizer to truncate
+    tokenizer = AutoTokenizer.from_pretrained(cfg.model, revision=cfg.revision)
+    tokenizer.model_max_length = min(tokenizer.model_max_length, cfg.token_batch_size)
+
+    # Do all the data loading and preprocessing on the main process
+    ds = load_data_string(cfg.data.dataset, cfg.data.split, streaming=cfg.streaming)
+
+    remove_columns = ds.column_names if cfg.drop_columns else None
+    ds = ds.map(
+        tokenize,
+        batched=True,
+        fn_kwargs=dict(args=cfg.data, tokenizer=tokenizer),
+        remove_columns=remove_columns,
+    )
+    if cfg.data.reward_column:
+        ds = ds.add_column(
+            "advantage",
+            estimate_advantage(ds, cfg),
+            new_fingerprint="advantage",  # type: ignore
+        )
+
+    world_size = torch.cuda.device_count()
+    if world_size <= 1:
+        # Run the worker directly if no distributed training is needed. This is great
+        # for debugging purposes.
+        worker(0, 1, cfg, ds)
+    else:
+        # Set up multiprocessing and distributed training
+        mp.set_sharing_strategy("file_system")
+
+        # Find an available port for distributed training
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            _, port = s.getsockname()
+
+        ctx = start_processes(
+            "build",
+            dist_worker,
+            args={i: (i, world_size, cfg, ds) for i in range(world_size)},
+            envs={
+                i: {
+                    "LOCAL_RANK": str(i),
+                    "MASTER_ADDR": "localhost",
+                    "MASTER_PORT": str(port),
+                }
+                for i in range(world_size)
+            },
+            logs_specs=DefaultLogsSpecs(),
+        )
+        ctx.wait()

bergson/collection.py

@@ -0,0 +1,214 @@
+import math
+from typing import Literal
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from datasets import Dataset, Value
+from tqdm.auto import tqdm
+from transformers import PreTrainedModel
+
+from .data import create_index, pad_and_tensor
+from .gradients import GradientCollector, GradientProcessor, HeadConfig
+from .peft import set_peft_enabled
+
+
+def collect_gradients(
+    model: PreTrainedModel,
+    data: Dataset,
+    processor: GradientProcessor,
+    path: str,
+    *,
+    batches: list[list[int]] | None = None,
+    kl_divergence: bool | None = None,
+    loss_reduction: Literal["mean", "sum"] = "mean",
+    skip_preconditioners: bool = False,
+    target_modules: set[str] | None = None,
+    head_cfgs: dict[str, HeadConfig] = {},
+):
+    """
+    Compute projected gradients using a subset of the dataset.
+    """
+    rank = dist.get_rank() if dist.is_initialized() else 0
+
+    # Batch size of one by default
+    if batches is None:
+        batches = [[idx] for idx in range(len(data))]
+
+    # Mutable state for the GradientCollector callback
+    mod_grads = {}
+    preconditioners = {}
+
+    # TODO: Handle this more elegantly
+    dtype = torch.float32 if model.dtype == torch.float32 else torch.float16
+    np_dtype = np.float32 if dtype == torch.float32 else np.float16
+    lo = torch.finfo(dtype).min
+    hi = torch.finfo(dtype).max
+
+    def callback(name: str, g: torch.Tensor):
+        g = g.flatten(1).clamp_(lo, hi)
+
+        # Asynchronously move the gradient to CPU and convert to fp16
+        mod_grads[name] = g.to(device="cpu", dtype=dtype, non_blocking=True)
+
+        # Compute the outer product of the flattened gradient
+        if not skip_preconditioners:
+            g = g.float()
+            preconditioner = preconditioners.get(name, None)
+            if preconditioner is None:
+                preconditioners[name] = g.mT @ g
+            else:
+                preconditioner.addmm_(g.mT, g)
+
+    collector = GradientCollector(
+        model.base_model,
+        callback,
+        processor,
+        target_modules=target_modules,
+        head_cfgs=head_cfgs,
+    )
+
+    # Allocate space ahead of time for the gradients
+    grad_sizes = {name: math.prod(s) for name, s in collector.shapes().items()}
+
+    # Allocate structured space ahead of time for the gradients
+    grad_buffer = create_index(
+        path, num_grads=len(data), grad_sizes=grad_sizes, dtype=np_dtype
+    )
+
+    per_doc_losses = torch.full(
+        (len(data),),
+        device=model.device,
+        dtype=dtype,
+        fill_value=0.0,
+    )
+
+    for indices in tqdm(batches, disable=rank != 0, desc="Building index"):
+        batch = data[indices]
+        x, y = pad_and_tensor(
+            batch["input_ids"],  # type: ignore
+            labels=batch.get("labels"),  # type: ignore
+            device=model.device,
+        )
+        masks = y[:, 1:] != -100
+        denoms = masks.sum(dim=1, dtype=dtype) if loss_reduction == "mean" else 1.0
+
+        if kl_divergence:
+            with torch.inference_mode():
+                set_peft_enabled(model, False)
+                ref_lps = torch.log_softmax(model(x).logits[:, :-1], dim=-1)
+                set_peft_enabled(model, True)
+
+            with collector:
+                ft_lps = torch.log_softmax(model(x).logits[:, :-1], dim=-1)
+
+                # Compute average KL across all unmasked tokens
+                kls = torch.sum(ft_lps.exp() * (ft_lps - ref_lps), dim=-1)
+                losses = torch.sum(kls * masks, dim=-1) / denoms
+                if "advantage" in batch:
+                    losses *= torch.tensor(batch["advantage"], device=losses.device)
+
+                losses.mean().backward()
+        else:
+            with collector:
+                logits = model(x).logits[:, :-1]
+
+                losses = F.cross_entropy(
+                    logits.reshape(-1, logits.size(-1)),
+                    y[:, 1:].flatten(),
+                    reduction="none",
+                ).reshape_as(y[:, 1:])
+                losses = losses.sum(1) / denoms
+                if "advantage" in batch:
+                    losses *= torch.tensor(batch["advantage"], device=losses.device)
+
+                losses.mean().backward()
+
+        # Weirdly you need to explicitly synchronize here in order to make sure that
+        # the nonblocking copies actually finish before we call .numpy()
+        model.zero_grad()
+        torch.cuda.synchronize()
+
+        # It turns out that it's very important for efficiency to write the gradients
+        # sequentially instead of first concatenating them, then writing to one vector
+        for module_name in mod_grads.keys():
+            grad_buffer[module_name][indices] = mod_grads[module_name].numpy()
+
+        mod_grads.clear()
+        per_doc_losses[indices] = losses.detach().type_as(per_doc_losses)
+
+    process_preconditioners(processor, preconditioners, len(data))
+
+    if dist.is_initialized():
+        dist.reduce(per_doc_losses, dst=0)
+
+    if rank == 0:
+        data = data.add_column(
+            "loss",
+            per_doc_losses.cpu().numpy(),
+            feature=Value("float16" if dtype == torch.float16 else "float32"),
+            new_fingerprint="loss",
+        )
+        data.save_to_disk(path + "/data.hf")
+
+        processor.save(path)
+
+    # Make sure the gradients are written to disk
+    grad_buffer.flush()
+
+
+def process_preconditioners(
+    processor: GradientProcessor,
+    preconditioners: dict[str, torch.Tensor],
+    len_data: int,
+):
+    """
+    Aggregate preconditioners across ranks and compute their eigen decomposition
+    distributed across all ranks.
+    """
+
+    rank = dist.get_rank() if dist.is_initialized() else 0
+    world_size = dist.get_world_size() if dist.is_initialized() else 1
+    preconditioners_eigen = {}
+    if rank == 0:
+        print("Saving preconditioners...")
+    for name, prec in preconditioners.items():
+        if dist.is_initialized():
+            dist.all_reduce(prec)
+
+        preconditioners[name] = prec / len_data
+
+    processor.preconditioners = preconditioners
+
+    if rank == 0:
+        print("Computing preconditioner eigen decompositions...")
+    names = list(preconditioners.keys())
+    names_per_rank = names[rank::world_size]
+
+    for name in names_per_rank:
+        original_dtype = preconditioners[name].dtype
+        prec = preconditioners[name].to(dtype=torch.float64)
+        eigvals, eigvecs = torch.linalg.eigh(prec)
+        preconditioners_eigen[name] = (
+            eigvals.to(dtype=original_dtype).contiguous(),
+            eigvecs.to(dtype=original_dtype).contiguous(),
+        )
+
+    if rank == 0:
+        print("Gathering and saving preconditioner eigen decompositions...")
+
+    for name in names:
+        prec = preconditioners[name]
+        if name not in preconditioners_eigen:
+            eigval = torch.zeros(prec.size(0), dtype=prec.dtype, device=prec.device)
+            eigvec = torch.zeros_like(prec)
+        else:
+            eigval, eigvec = preconditioners_eigen[name]
+
+        dist.all_reduce(eigval, op=dist.ReduceOp.SUM) if dist.is_initialized() else None
+        dist.all_reduce(eigvec, op=dist.ReduceOp.SUM) if dist.is_initialized() else None
+
+        preconditioners_eigen[name] = (eigval, eigvec)
+    if rank == 0:
+        processor.preconditioners_eigen = preconditioners_eigen