Upload 11 files

LICENSE

@@ -0,0 +1,17 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+Copyright 2026 Angshul Majumdar
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

README.md

@@ -1,3 +1,100 @@
 ---
+library_name: pytorch
 license: apache-2.0
+tags:
+  - pruning
+  - sparsity
+  - optimization
+  - higher-order
+  - transformers
 ---
+
+# Higher-Order Pruner
+
+A reusable Python pruning method hosted in a Hugging Face **model repository**.
+This repository is not a pruned checkpoint and does not contain benchmark
+results. It supplies a direct code API for pruning compatible unquantized
+PyTorch models.
+
+The method combines calibration-derived KFAC-diagonal second-order information
+with a clipped diagonal higher-order finite-difference estimate, exact global
+sparsity allocation, and monotone held-out refinement. It does not initialize
+from Wanda, SparseGPT, magnitude pruning, or a previously pruned checkpoint.
+
+## Install from this model repository
+
+```bash
+pip install "git+https://huggingface.co/OWNER/REPOSITORY"
+```
+
+Then load the method configuration from the Hub and use it in code:
+
+```python
+from higher_order_pruner import HigherOrderPruner
+
+pruner = HigherOrderPruner.from_pretrained("OWNER/REPOSITORY")
+
+result = pruner.prune_hf_model(
+    "Qwen/Qwen3-4B-Base",
+    calibration_texts=my_calibration_texts,
+    validation_texts=my_heldout_texts,
+    sparsity=0.50,
+    output_dir="./pruned-model",
+    sequence_length=512,
+    device="cuda",
+)
+print(result["report"]["achieved_sparsity"])
+```
+
+`pruned-model/checkpoint` is a normal Hugging Face checkpoint saved with
+`save_pretrained` and safetensors.
+
+## Generic PyTorch API
+
+The core is not restricted to language models. Supply any differentiable
+scalar loss and batches already placed on the desired device:
+
+```python
+from higher_order_pruner import HigherOrderPruner
+
+pruner = HigherOrderPruner.from_pretrained("OWNER/REPOSITORY")
+
+def loss_fn(model, batch):
+    x, target = batch
+    return (model(x) - target).square().mean()
+
+report = pruner.prune(
+    model,
+    calibration_batches,
+    loss_fn,
+    sparsity=0.50,
+    validation_batches=heldout_batches,
+    progress_dir="./pruning-progress",
+)
+```
+
+For nonstandard architectures, pass explicit `module_names=[...]`. The method
+prunes unique floating-point `torch.nn.Linear.weight` tensors. It refuses
+quantized weights and does not promise that arbitrary custom kernels, tied
+non-Linear tensors, or remote-code-only architectures are compatible.
+
+## Important behavior
+
+- Calibration statistics and higher-order estimates are computed from the
+  caller's model and data at runtime; no Qwen/Falcon statistics are embedded.
+- The requested zero count is enforced exactly across selected weights.
+- Existing exact zeros are included in that budget.
+- The HF adapter auto-tunes its measured curvature kernel for throughput under
+  a 95% VRAM ceiling. It does not require a specific GPU model.
+- Unstructured zeros require an appropriate sparse inference kernel to produce
+  inference speedups.
+
+## Local development
+
+```bash
+pip install -e ".[test]"
+pytest -q
+```
+
+`upload_to_hub.py` creates a new Hugging Face **model** repository and uploads
+this package. It refuses to overwrite an existing repository.

higher_order_pruner/__init__.py

@@ -0,0 +1,17 @@
+"""Model-agnostic higher-order pruning for PyTorch linear layers."""
+
+from .core import PruneConfig, PruneReport, discover_linear_groups, prune
+from .hf import causal_lm_loss, load_causal_lm, prepare_texts, prune_hf_model
+from .pruner import HigherOrderPruner
+
+__all__ = [
+    "HigherOrderPruner",
+    "PruneConfig",
+    "PruneReport",
+    "discover_linear_groups",
+    "prune",
+    "prune_hf_model",
+    "load_causal_lm",
+    "prepare_texts",
+    "causal_lm_loss",
+]

higher_order_pruner/core.py

@@ -0,0 +1,421 @@
+"""Higher-order, KFAC-diagonal, monotone layerwise pruning.
+
+No model names, transformer attributes, dataset names, or evaluation tasks are
+assumed here. The caller owns the loss and supplies calibration batches.
+"""
+
+from __future__ import annotations
+
+import gc
+import json
+import math
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any, Callable, Iterable, Sequence
+
+import torch
+from torch import nn
+
+
+LossFn = Callable[[nn.Module, Any], torch.Tensor]
+Group = list[tuple[str, nn.Linear]]
+
+
+@dataclass(frozen=True)
+class PruneConfig:
+    sparsity: float
+    probes: int = 2
+    finite_difference_relative_step: float = 0.02
+    cubic_cap: float = 0.5
+    damping_ratio: float = 0.01
+    trust_ratio: float = 0.01
+    backtracking: tuple[float, ...] = (1.0, 0.5, 0.25, 0.125, 0.0625)
+    seed: int = 91000
+
+    def __post_init__(self) -> None:
+        if not 0 <= self.sparsity < 1:
+            raise ValueError("sparsity must lie in [0, 1)")
+        if self.probes < 1:
+            raise ValueError("probes must be positive")
+        if self.finite_difference_relative_step <= 0:
+            raise ValueError("finite_difference_relative_step must be positive")
+        if not 0 <= self.cubic_cap <= 1:
+            raise ValueError("cubic_cap must lie in [0, 1]")
+        if self.damping_ratio <= 0 or self.trust_ratio <= 0:
+            raise ValueError("damping_ratio and trust_ratio must be positive")
+        if not self.backtracking or any(a <= 0 for a in self.backtracking):
+            raise ValueError("backtracking must contain positive step sizes")
+
+
+@dataclass
+class PruneReport:
+    target_sparsity: float
+    achieved_sparsity: float
+    prunable_weights: int
+    exact_zero_weights: int
+    seconds: float
+    layers: list[dict[str, Any]]
+
+
+def discover_linear_groups(model: nn.Module, *, module_names: Sequence[str] | None = None) -> list[Group]:
+    """Prefer the model's repeated block stack; otherwise group linears singly.
+
+    module_names may explicitly select nn.Linear module paths. This is the
+    adapter for unusual architectures (including MoE expert placement).
+    """
+    named = dict(model.named_modules())
+    if module_names is not None:
+        selected = []
+        for name in module_names:
+            if name not in named or not isinstance(named[name], nn.Linear):
+                raise ValueError(f"Not a Linear module: {name}")
+            selected.append((name, named[name]))
+        if len({name for name, _ in selected}) != len(selected):
+            raise ValueError("module_names contains duplicates")
+        if not selected:
+            raise ValueError("No Linear modules selected")
+        return [[item] for item in selected]
+
+    candidates: list[tuple[int, int, str, nn.ModuleList]] = []
+    for name, module in named.items():
+        if isinstance(module, nn.ModuleList) and len(module) >= 2:
+            counts = [sum(isinstance(x, nn.Linear) for x in block.modules()) for block in module]
+            if min(counts) > 0:
+                candidates.append((sum(counts), len(module), name, module))
+    if candidates:
+        _, _, stack_name, stack = max(candidates, key=lambda x: (x[0], x[1]))
+        groups: list[Group] = []
+        for i, block in enumerate(stack):
+            prefix = f"{stack_name}.{i}" if stack_name else str(i)
+            groups.append([(f"{prefix}.{name}" if name else prefix, layer)
+                           for name, layer in block.named_modules() if isinstance(layer, nn.Linear)])
+    else:
+        groups = [[(name, layer)] for name, layer in named.items() if isinstance(layer, nn.Linear)]
+
+    seen: set[int] = set()
+    unique: list[Group] = []
+    for group in groups:
+        items = []
+        for name, layer in group:
+            ident = id(layer.weight)
+            if ident not in seen:
+                seen.add(ident)
+                items.append((name, layer))
+        if items:
+            unique.append(items)
+    if not unique:
+        raise ValueError("No unique nn.Linear weights found; select compatible modules explicitly")
+    return unique
+
+
+def _loss(model: nn.Module, batch: Any, loss_fn: LossFn) -> torch.Tensor:
+    value = loss_fn(model, batch)
+    if not isinstance(value, torch.Tensor) or value.ndim != 0 or not torch.isfinite(value).item():
+        raise ValueError("loss_fn must return a finite scalar torch.Tensor")
+    return value
+
+
+def _mean_loss(model: nn.Module, batches: Sequence[Any], loss_fn: LossFn) -> float:
+    with torch.no_grad():
+        return sum(float(_loss(model, b, loss_fn).detach()) for b in batches) / len(batches)
+
+
+def _gradient(model: nn.Module, group: Group, batches: Sequence[Any], loss_fn: LossFn,
+              base: dict[str, torch.Tensor], u: dict[str, torch.Tensor] | None = None,
+              v: dict[str, torch.Tensor] | None = None, us: int = 0, vs: int = 0,
+              steps: dict[str, float] | None = None) -> dict[str, torch.Tensor]:
+    with torch.no_grad():
+        for name, layer in group:
+            if u is not None and v is not None and steps is not None:
+                candidate = base[name].float() + steps[name] * (us * u[name].float() + vs * v[name].float())
+                layer.weight.copy_(candidate.to(layer.weight.dtype))
+            layer.weight.requires_grad_(True)
+    sums = {name: torch.zeros_like(layer.weight, dtype=torch.float32) for name, layer in group}
+    try:
+        for batch in batches:
+            model.zero_grad(set_to_none=True)
+            _loss(model, batch, loss_fn).backward()
+            for name, layer in group:
+                if layer.weight.grad is None:
+                    raise RuntimeError(f"Selected module has no loss gradient: {name}")
+                sums[name].add_(layer.weight.grad.detach().float())
+        return {name: value / len(batches) for name, value in sums.items()}
+    finally:
+        model.zero_grad(set_to_none=True)
+        with torch.no_grad():
+            for name, layer in group:
+                layer.weight.copy_(base[name])
+                layer.weight.requires_grad_(False)
+
+
+def _kfac_diagonal(model: nn.Module, group: Group, batches: Sequence[Any],
+                   loss_fn: LossFn) -> dict[str, tuple[torch.Tensor, torch.Tensor]]:
+    """Collect E[output-gradient²] and E[input²] on the current model."""
+    stats: dict[str, dict[str, Any]] = {
+        name: {"a": torch.zeros(layer.in_features, device=layer.weight.device, dtype=torch.float32),
+               "g": torch.zeros(layer.out_features, device=layer.weight.device, dtype=torch.float32),
+               "na": 0, "ng": 0}
+        for name, layer in group
+    }
+    hooks = []
+
+    def make_hook(name: str):
+        def forward_hook(_module: nn.Module, inputs: tuple[Any, ...], output: torch.Tensor) -> None:
+            x = inputs[0].detach().reshape(-1, stats[name]["a"].numel()).float()
+            stats[name]["a"].add_(x.square().sum(dim=0))
+            stats[name]["na"] += x.shape[0]
+
+            def backward_hook(grad: torch.Tensor) -> None:
+                y = grad.detach().reshape(-1, stats[name]["g"].numel()).float()
+                stats[name]["g"].add_(y.square().sum(dim=0))
+                stats[name]["ng"] += y.shape[0]
+
+            output.register_hook(backward_hook)
+        return forward_hook
+
+    try:
+        for name, layer in group:
+            layer.weight.requires_grad_(True)
+            hooks.append(layer.register_forward_hook(make_hook(name)))
+        for batch in batches:
+            model.zero_grad(set_to_none=True)
+            _loss(model, batch, loss_fn).backward()
+        answer = {}
+        for name, item in stats.items():
+            if not item["na"] or not item["ng"]:
+                raise RuntimeError(f"No activation/gradient statistics for {name}")
+            answer[name] = (item["g"] / item["ng"], item["a"] / item["na"])
+        return answer
+    finally:
+        for hook in hooks:
+            hook.remove()
+        model.zero_grad(set_to_none=True)
+        for _, layer in group:
+            layer.weight.requires_grad_(False)
+
+
+def _higher_order_diagonal(model: nn.Module, group: Group, batches: Sequence[Any], loss_fn: LossFn,
+                           cfg: PruneConfig, group_index: int) -> dict[str, torch.Tensor]:
+    base = {name: layer.weight.detach().clone() for name, layer in group}
+    estimate = {name: torch.zeros_like(layer.weight, dtype=torch.float32) for name, layer in group}
+    gen = torch.Generator(device=group[0][1].weight.device)
+    gen.manual_seed(cfg.seed + group_index)
+    for _ in range(cfg.probes):
+        u = {name: torch.randint(0, 2, layer.weight.shape, device=layer.weight.device,
+                                 dtype=torch.int8, generator=gen).mul_(2).sub_(1)
+             for name, layer in group}
+        v = {name: torch.randint(0, 2, layer.weight.shape, device=layer.weight.device,
+                                 dtype=torch.int8, generator=gen).mul_(2).sub_(1)
+             for name, layer in group}
+        steps = {name: max(cfg.finite_difference_relative_step *
+                           base[name].float().square().mean().sqrt().item(), 1e-6)
+                 for name, _ in group}
+        for us, vs, sign in ((1, 1, 1), (1, -1, -1), (-1, 1, -1), (-1, -1, 1)):
+            grads = _gradient(model, group, batches, loss_fn, base, u, v, us, vs, steps)
+            for name in grads:
+                estimate[name].add_(grads[name] * u[name].float() * v[name].float(),
+                                    alpha=sign / (4 * steps[name] ** 2 * cfg.probes))
+            del grads
+        del u, v
+    for name in estimate:
+        sample = estimate[name].reshape(-1)[::max(1, estimate[name].numel() // 131072)]
+        limit = 10 * sample.abs().median().clamp_min(1e-12)
+        estimate[name].clamp_(-limit, limit)
+    return estimate
+
+
+def _budgets(groups: list[Group], sparsity: float) -> list[int]:
+    sizes = [sum(layer.weight.numel() for _, layer in group) for group in groups]
+    wanted = round(sum(sizes) * sparsity)
+    existing = [sum(int((layer.weight.detach() == 0).sum().item())
+                    for _, layer in group) for group in groups]
+    if sum(existing) > wanted:
+        raise ValueError(
+            f"The selected weights already contain {sum(existing)} exact zeros, "
+            f"which exceeds the requested global budget {wanted}."
+        )
+    exact = [size * sparsity for size in sizes]
+    quotas = [max(existing[i], math.floor(exact[i])) for i in range(len(groups))]
+    if sum(quotas) > wanted:
+        # A group with many pre-existing zeros borrows budget from groups that
+        # are still above their immutable lower bound.
+        order = sorted(range(len(groups)), key=lambda i: (quotas[i] - existing[i], i), reverse=True)
+        remaining = sum(quotas) - wanted
+        for i in order:
+            take = min(remaining, quotas[i] - existing[i])
+            quotas[i] -= take
+            remaining -= take
+            if remaining == 0:
+                break
+    else:
+        order = sorted(range(len(groups)), key=lambda i: (exact[i] - math.floor(exact[i]), -i), reverse=True)
+        remaining = wanted - sum(quotas)
+        for i in order:
+            if remaining == 0:
+                break
+            if quotas[i] < sizes[i]:
+                quotas[i] += 1
+                remaining -= 1
+        # For unusual pre-existing-zero distributions, the fractional pass
+        # may not consume the whole budget. Fill deterministically.
+        if remaining:
+            for i in range(len(groups)):
+                take = min(remaining, sizes[i] - quotas[i])
+                quotas[i] += take
+                remaining -= take
+                if remaining == 0:
+                    break
+    if sum(quotas) != wanted:
+        raise RuntimeError("Could not distribute the exact global sparsity budget")
+    return quotas
+
+
+def _mask_group(group: Group, factors: dict[str, tuple[torch.Tensor, torch.Tensor]],
+                higher_order: dict[str, torch.Tensor], quota: int, cfg: PruneConfig):
+    scores = []
+    existing = []
+    sizes = []
+    for name, layer in group:
+        w = layer.weight.detach().float()
+        g, a = factors[name]
+        h = g[:, None] * a[None, :]
+        quadratic = 0.5 * h * w.square()
+        cubic_raw = -higher_order[name] * w.pow(3) / 6
+        cubic_limit = cfg.cubic_cap * quadratic.abs()
+        cubic = torch.maximum(torch.minimum(cubic_raw, cubic_limit), -cubic_limit)
+        score = (quadratic + cubic).clamp_min_(0).reshape(-1)
+        score = torch.nan_to_num(score, nan=float("inf"), posinf=float("inf"))
+        scores.append(score)
+        existing.append((layer.weight.detach().reshape(-1) == 0))
+        sizes.append(layer.weight.numel())
+    joint = torch.cat(scores)
+    already = torch.cat(existing)
+    n_zero = int(already.sum().item())
+    if quota < n_zero:
+        raise ValueError(f"Target quota {quota} is below {n_zero} pre-existing exact zeros")
+    additional = quota - n_zero
+    available = torch.nonzero(~already).flatten()
+    if additional > available.numel():
+        raise RuntimeError("Insufficient nonzero weights to meet the sparsity target")
+    chosen = available[torch.topk(joint[available], k=additional, largest=False, sorted=False).indices] \
+        if additional else available[:0]
+    mask = already.clone()
+    mask[chosen] = True
+    support = {}
+    base = {}
+    offset = 0
+    with torch.no_grad():
+        for (name, layer), length in zip(group, sizes):
+            zero = mask[offset:offset + length].reshape_as(layer.weight)
+            layer.weight[zero] = 0
+            support[name] = ~zero
+            base[name] = layer.weight.detach().clone()
+            offset += length
+    return support, base
+
+
+def _refine(model: nn.Module, group: Group, factors: dict[str, tuple[torch.Tensor, torch.Tensor]],
+            higher_order: dict[str, torch.Tensor], support: dict[str, torch.Tensor],
+            base: dict[str, torch.Tensor], train: Sequence[Any], validation: Sequence[Any],
+            loss_fn: LossFn, cfg: PruneConfig) -> tuple[float, float, float]:
+    before = _mean_loss(model, validation, loss_fn)
+    gradients = _gradient(model, group, train, loss_fn, base)
+    directions = {}
+    for name, layer in group:
+        g, a = factors[name]
+        h = g[:, None] * a[None, :]
+        h = h + cfg.damping_ratio * h.mean().clamp_min(1e-12)
+        t = higher_order[name]
+        grad = gradients[name]
+        discriminant = h.square() - 2 * t * grad
+        root = discriminant.clamp_min(0).sqrt()
+        linear_step = -grad / h
+        cubic_step = -2 * grad / (h + root).clamp_min(1e-20)
+        direction = torch.where((discriminant > 0) & (t.abs() > 1e-12) &
+                                torch.isfinite(cubic_step), cubic_step, linear_step)
+        direction.masked_fill_(~support[name], 0)
+        radius = cfg.trust_ratio * layer.weight.detach().float()[support[name]].norm()
+        norm = direction[support[name]].norm()
+        if norm > radius:
+            direction.mul_(radius / norm.clamp_min(1e-12))
+        directions[name] = direction
+    accepted = 0.0
+    after = before
+    with torch.no_grad():
+        for alpha in cfg.backtracking:
+            for name, layer in group:
+                candidate = base[name].float() + alpha * directions[name]
+                candidate.masked_fill_(~support[name], 0)
+                layer.weight.copy_(candidate.to(layer.weight.dtype))
+            loss = _mean_loss(model, validation, loss_fn)
+            if loss <= before + 1e-5:
+                accepted, after = alpha, loss
+                break
+        if accepted == 0:
+            for name, layer in group:
+                layer.weight.copy_(base[name])
+    return before, after, accepted
+
+
+def prune(model: nn.Module, calibration_batches: Iterable[Any], loss_fn: LossFn,
+          config: PruneConfig, *, validation_batches: Iterable[Any] | None = None,
+          module_names: Sequence[str] | None = None, progress_dir: str | Path | None = None,
+          on_progress: Callable[[dict[str, Any]], None] | None = None) -> PruneReport:
+    """Prune an arbitrary compatible torch model in place.
+
+    The caller defines the forward/loss via loss_fn. Only nn.Linear weights
+    are pruned; tensors and batches may live on CPU or GPU. The model must fit
+    on the chosen device for forward/backward and may not be quantized.
+    """
+    train = list(calibration_batches)
+    validation = list(validation_batches) if validation_batches is not None else train[-1:]
+    if not train or not validation:
+        raise ValueError("Nonempty calibration and validation batches required")
+    groups = discover_linear_groups(model, module_names=module_names)
+    quotas = _budgets(groups, config.sparsity)
+    total = sum(layer.weight.numel() for group in groups for _, layer in group)
+    if any(not layer.weight.is_floating_point() for group in groups for _, layer in group):
+        raise ValueError("Pruning requires unquantized floating point nn.Linear weights")
+    old_training = model.training
+    original_flags = {id(p): p.requires_grad for p in model.parameters()}
+    model.eval()
+    model.requires_grad_(False)
+    report_rows = []
+    started = time.time()
+    dest = Path(progress_dir) if progress_dir is not None else None
+    if dest is not None:
+        dest.mkdir(parents=True, exist_ok=True)
+    try:
+        for i, (group, quota) in enumerate(zip(groups, quotas)):
+            layer_start = time.time()
+            factors = _kfac_diagonal(model, group, train, loss_fn)
+            higher_order = _higher_order_diagonal(model, group, train, loss_fn, config, i)
+            support, base = _mask_group(group, factors, higher_order, quota, config)
+            before, after, alpha = _refine(model, group, factors, higher_order, support, base,
+                                           train, validation, loss_fn, config)
+            row = {"group": i, "modules": [name for name, _ in group], "zeros": quota,
+                   "loss_before": before, "loss_after": after, "accepted_step": alpha,
+                   "seconds": time.time() - layer_start}
+            report_rows.append(row)
+            if dest is not None:
+                from safetensors.torch import save_file
+                save_file({name: layer.weight.detach().cpu().contiguous()
+                           for name, layer in group}, str(dest / f"group_{i:04d}.safetensors"))
+                (dest / "progress.json").write_text(json.dumps({"groups_completed": i + 1,
+                    "total_groups": len(groups), "layers": report_rows}, indent=2))
+            if on_progress is not None:
+                on_progress(row)
+            del factors, higher_order, support, base
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        actual = sum(int((layer.weight == 0).sum().item()) for group in groups for _, layer in group)
+        if actual != round(total * config.sparsity):
+            raise RuntimeError(f"Exact zero count mismatch: {actual} != {round(total * config.sparsity)}")
+        return PruneReport(config.sparsity, actual / total, total, actual,
+                           time.time() - started, report_rows)
+    finally:
+        model.train(old_training)
+        for parameter in model.parameters():
+            parameter.requires_grad_(original_flags[id(parameter)])

higher_order_pruner/hf.py

@@ -0,0 +1,166 @@
+"""Hugging Face CausalLM adapter for the generic pruning core."""
+
+from __future__ import annotations
+
+import gc
+import json
+import time
+from pathlib import Path
+from typing import Callable, Sequence
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from .core import PruneConfig, _higher_order_diagonal, _kfac_diagonal, discover_linear_groups, prune
+
+
+def load_causal_lm(model_id: str, *, revision: str = "main", device: str = "cuda"):
+    """Load an ordinary unquantized HF causal LM without remote Python code."""
+    if device.startswith("cuda") and not torch.cuda.is_available():
+        raise RuntimeError("A CUDA GPU is required for this requested device")
+    tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision,
+                                               trust_remote_code=False, use_fast=True)
+    if tokenizer.pad_token_id is None:
+        if tokenizer.eos_token_id is None:
+            raise ValueError("The tokenizer needs a padding or EOS token")
+        tokenizer.pad_token = tokenizer.eos_token
+    dtype = torch.bfloat16 if device.startswith("cuda") and torch.cuda.is_bf16_supported() \
+        else (torch.float16 if device.startswith("cuda") else torch.float32)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id, revision=revision, dtype=dtype, trust_remote_code=False,
+        low_cpu_mem_usage=True, device_map={"": device}, use_safetensors=True,
+    )
+    if getattr(model, "is_quantized", False) or getattr(model, "hf_quantizer", None) is not None:
+        raise ValueError("This method requires an unquantized model")
+    model.eval()
+    if hasattr(model.config, "use_cache"):
+        model.config.use_cache = False
+    return model, tokenizer
+
+
+def prepare_texts(tokenizer, texts: Sequence[str], *, sequence_length: int = 256,
+                  batch_size: int = 1, device: str = "cuda") -> list[dict[str, torch.Tensor]]:
+    if not texts or not all(isinstance(text, str) and text.strip() for text in texts):
+        raise ValueError("Supply nonempty calibration text strings")
+    if not 16 <= sequence_length <= 2048:
+        raise ValueError("sequence_length must lie between 16 and 2048")
+    if batch_size < 1:
+        raise ValueError("batch_size must be positive")
+    outputs = []
+    for offset in range(0, len(texts), batch_size):
+        batch = tokenizer(list(texts[offset:offset + batch_size]),
+                          padding=True, truncation=True, max_length=sequence_length,
+                          return_tensors="pt")
+        ids = batch["input_ids"]
+        mask = batch["attention_mask"]
+        labels = ids.clone()
+        labels[mask == 0] = -100
+        if (labels[:, 1:] != -100).sum().item() == 0:
+            raise ValueError("Calibration sample has fewer than two usable tokens")
+        outputs.append({"input_ids": ids.to(device), "attention_mask": mask.to(device),
+                        "labels": labels.to(device)})
+    return outputs
+
+
+def causal_lm_loss(model, batch: dict[str, torch.Tensor]) -> torch.Tensor:
+    return model(**batch, use_cache=False).loss
+
+
+def tune_batch(model, tokenizer, texts: Sequence[str], *, sequence_length: int,
+               device: str, hard_fraction: float = 0.95,
+               candidates: tuple[int, ...] = (1, 2, 4, 8, 12, 16)) -> dict:
+    """Measure actual KFAC and higher-order gradient kernels, not just a forward."""
+    if not device.startswith("cuda"):
+        return {"batch_size": 1, "measurements": [], "device": device}
+    groups = discover_linear_groups(model)
+    representative = max(groups, key=lambda group: sum(m.weight.numel() for _, m in group))
+    total = torch.cuda.get_device_properties(0).total_memory
+    measurements = []
+    for batch_size in candidates:
+        if len(texts) < batch_size:
+            continue
+        payload = None
+        try:
+            payload = prepare_texts(tokenizer, texts[:batch_size],
+                                    sequence_length=sequence_length,
+                                    batch_size=batch_size, device=device)
+            torch.cuda.empty_cache()
+            torch.cuda.reset_peak_memory_stats()
+            t0 = time.monotonic()
+            factors = _kfac_diagonal(model, representative, payload, causal_lm_loss)
+            # One Rademacher probe uses the same peak allocations as a
+            # production probe without repeating the full calibration corpus.
+            probe = _higher_order_diagonal(model, representative, payload, causal_lm_loss,
+                                           PruneConfig(sparsity=0.5, probes=1), 0)
+            torch.cuda.synchronize()
+            seconds = time.monotonic() - t0
+            peak = torch.cuda.max_memory_reserved()
+            measurements.append({"batch_size": batch_size, "seconds": seconds,
+                                 "peak_gib": peak / 2**30,
+                                 "safe": peak <= total * hard_fraction})
+            del factors, probe
+        except torch.OutOfMemoryError:
+            measurements.append({"batch_size": batch_size, "oom": True, "safe": False})
+            model.zero_grad(set_to_none=True)
+        finally:
+            del payload
+            gc.collect()
+            torch.cuda.empty_cache()
+    safe = [m for m in measurements if m.get("safe")]
+    if not safe:
+        raise RuntimeError("No safe batch size found; use a smaller model or shorter sequences")
+    # Prefer throughput; memory occupancy is not an optimization objective.
+    best = max(safe, key=lambda m: m["batch_size"] / m["seconds"])
+    return {"batch_size": best["batch_size"], "measurements": measurements,
+            "device": device, "sequence_length": sequence_length}
+
+
+def prune_hf_model(model_id: str, calibration_texts: Sequence[str], *,
+                   validation_texts: Sequence[str] | None = None, sparsity: float,
+                   output_dir: str | Path, revision: str = "main",
+                   sequence_length: int = 256, probes: int = 2,
+                   module_names: Sequence[str] | None = None, device: str = "cuda",
+                   on_progress: Callable[[dict], None] | None = None) -> dict:
+    """Download, calibrate, prune, and save any supported HF causal LM.
+
+    'Supported' means that the model loads with AutoModelForCausalLM without
+    remote code, has nn.Linear modules and a differentiable language-model
+    loss, and fits with backward activations on the selected device.
+    """
+    dest = Path(output_dir)
+    dest.mkdir(parents=True, exist_ok=True)
+    model, tokenizer = load_causal_lm(model_id, revision=revision, device=device)
+    try:
+        cfg = PruneConfig(sparsity=sparsity, probes=probes)
+        if module_names is None:
+            tuned = tune_batch(model, tokenizer, calibration_texts,
+                               sequence_length=sequence_length, device=device)
+        else:
+            # A custom architecture selection can be much larger than a block;
+            # use a conservative batch rather than an unrelated probe.
+            tuned = {"batch_size": 1, "measurements": [], "custom_selection": True}
+        batch_size = tuned["batch_size"]
+        train = prepare_texts(tokenizer, calibration_texts, sequence_length=sequence_length,
+                              batch_size=batch_size, device=device)
+        validation = prepare_texts(tokenizer, validation_texts or calibration_texts[-2:],
+                                   sequence_length=sequence_length, batch_size=1, device=device)
+        try:
+            report = prune(model, train, causal_lm_loss, cfg, validation_batches=validation,
+                           module_names=module_names, progress_dir=dest / "progress",
+                           on_progress=on_progress)
+        finally:
+            del train, validation
+        checkpoint = dest / "checkpoint"
+        model.save_pretrained(checkpoint, safe_serialization=True, max_shard_size="2GB")
+        tokenizer.save_pretrained(checkpoint)
+        result = {"status": "complete", "source_model": model_id, "revision": revision,
+                  "checkpoint": str(checkpoint), "sparsity": sparsity,
+                  "report": vars(report), "batch_tuning": tuned}
+        (dest / "report.json").write_text(json.dumps(result, indent=2))
+        return result
+    finally:
+        model.cpu()
+        del model
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()

higher_order_pruner/pruner.py

@@ -0,0 +1,118 @@
+"""Public, reusable pruning interface for Hugging Face model repositories."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, replace
+from pathlib import Path
+from typing import Any, Callable, Iterable, Sequence
+
+from huggingface_hub import hf_hub_download
+from torch import nn
+
+from .core import PruneConfig, PruneReport, prune
+from .hf import prune_hf_model
+
+
+class HigherOrderPruner:
+    """Configuration-backed API for higher-order pruning.
+
+    The object contains method defaults, not model weights. ``from_pretrained``
+    reads those defaults from a Hugging Face *model* repository. The same
+    instance can prune any compatible model; it is not tied to Qwen or Falcon.
+    """
+
+    CONFIG_NAME = "pruner_config.json"
+
+    def __init__(self, config: PruneConfig | None = None):
+        self.config = config or PruneConfig(sparsity=0.5)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        repo_id: str,
+        *,
+        revision: str = "main",
+        token: str | bool | None = None,
+        cache_dir: str | Path | None = None,
+        local_files_only: bool = False,
+    ) -> "HigherOrderPruner":
+        path = hf_hub_download(
+            repo_id=repo_id,
+            filename=cls.CONFIG_NAME,
+            revision=revision,
+            token=token,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+        payload = json.loads(Path(path).read_text())
+        if payload.get("format") != "higher-order-pruner-v1":
+            raise ValueError("Unsupported pruner configuration format")
+        values = dict(payload["defaults"])
+        if "backtracking" in values:
+            values["backtracking"] = tuple(values["backtracking"])
+        return cls(PruneConfig(**values))
+
+    def save_pretrained(self, directory: str | Path) -> Path:
+        destination = Path(directory)
+        destination.mkdir(parents=True, exist_ok=True)
+        path = destination / self.CONFIG_NAME
+        path.write_text(json.dumps({
+            "format": "higher-order-pruner-v1",
+            "method": "diagonal-higher-order-kfac",
+            "defaults": asdict(self.config),
+        }, indent=2) + "\n")
+        return path
+
+    def prune(
+        self,
+        model: nn.Module,
+        calibration_batches: Iterable[Any],
+        loss_fn: Callable[[nn.Module, Any], Any],
+        *,
+        sparsity: float | None = None,
+        validation_batches: Iterable[Any] | None = None,
+        module_names: Sequence[str] | None = None,
+        progress_dir: str | Path | None = None,
+        on_progress: Callable[[dict[str, Any]], None] | None = None,
+    ) -> PruneReport:
+        config = self.config if sparsity is None else replace(self.config, sparsity=sparsity)
+        return prune(
+            model,
+            calibration_batches,
+            loss_fn,
+            config,
+            validation_batches=validation_batches,
+            module_names=module_names,
+            progress_dir=progress_dir,
+            on_progress=on_progress,
+        )
+
+    def prune_hf_model(
+        self,
+        model_id: str,
+        calibration_texts: Sequence[str],
+        *,
+        output_dir: str | Path,
+        sparsity: float | None = None,
+        validation_texts: Sequence[str] | None = None,
+        revision: str = "main",
+        sequence_length: int = 256,
+        module_names: Sequence[str] | None = None,
+        device: str = "cuda",
+        on_progress: Callable[[dict], None] | None = None,
+    ) -> dict:
+        selected = self.config.sparsity if sparsity is None else sparsity
+        return prune_hf_model(
+            model_id,
+            calibration_texts,
+            validation_texts=validation_texts,
+            sparsity=selected,
+            output_dir=output_dir,
+            revision=revision,
+            sequence_length=sequence_length,
+            probes=self.config.probes,
+            module_names=module_names,
+            device=device,
+            on_progress=on_progress,
+        )

pruner_config.json

@@ -0,0 +1,14 @@
+{
+  "format": "higher-order-pruner-v1",
+  "method": "diagonal-higher-order-kfac",
+  "defaults": {
+    "sparsity": 0.5,
+    "probes": 2,
+    "finite_difference_relative_step": 0.02,
+    "cubic_cap": 0.5,
+    "damping_ratio": 0.01,
+    "trust_ratio": 0.01,
+    "backtracking": [1.0, 0.5, 0.25, 0.125, 0.0625],
+    "seed": 91000
+  }
+}

pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["setuptools>=70", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "higher-order-pruner"
+version = "0.1.0"
+description = "Model-agnostic higher-order pruning for PyTorch and Hugging Face models"
+readme = "README.md"
+requires-python = ">=3.10"
+license = {text = "Apache-2.0"}
+authors = [{name = "Angshul Majumdar"}]
+dependencies = [
+  "torch>=2.4,<3",
+  "transformers>=4.57,<5",
+  "accelerate>=1.2,<2",
+  "safetensors>=0.5,<1",
+  "huggingface_hub>=0.34,<2",
+]
+
+[project.optional-dependencies]
+test = ["pytest>=8,<10"]
+
+[tool.setuptools.packages.find]
+include = ["higher_order_pruner*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.4,<3
+transformers>=4.57,<5
+accelerate>=1.2,<2
+safetensors>=0.5,<1
+huggingface_hub>=0.34,<2

tests/test_core.py

@@ -0,0 +1,33 @@
+"""Tiny CPU smoke test; run with pytest after installing requirements."""
+
+import torch
+from torch import nn
+
+from higher_order_pruner import PruneConfig, discover_linear_groups, prune
+
+
+class TinyNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layers = nn.ModuleList([nn.Linear(5, 5), nn.Linear(5, 5)])
+
+    def forward(self, x):
+        return self.layers[1](self.layers[0](x).tanh())
+
+
+def test_exact_budget_and_general_loss():
+    torch.manual_seed(17)
+    model = TinyNet()
+    batches = [(torch.randn(4, 5), torch.randn(4, 5)) for _ in range(3)]
+
+    def mse(m, batch):
+        x, y = batch
+        return (m(x) - y).square().mean()
+
+    assert len(discover_linear_groups(model)) == 2
+    report = prune(model, batches[:2], mse, PruneConfig(sparsity=0.3, probes=1),
+                   validation_batches=batches[2:])
+    assert report.prunable_weights == 50
+    assert report.exact_zero_weights == 15
+    assert sum(int((p.weight == 0).sum()) for p in model.layers) == 15
+    assert all(row["loss_after"] <= row["loss_before"] + 1e-5 for row in report.layers)

upload_to_hub.py

@@ -0,0 +1,36 @@
+"""Create a new Hugging Face MODEL repository for this reusable method."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+from huggingface_hub import HfApi
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("repo_id", help="For example: Angshul/higher-order-pruner")
+    parser.add_argument("--public", action="store_true", help="Create a public repository")
+    args = parser.parse_args()
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        raise SystemExit("Set a write-scoped HF_TOKEN environment variable")
+    api = HfApi(token=token)
+    if api.repo_exists(args.repo_id, repo_type="model"):
+        raise SystemExit(f"Refusing to overwrite existing model repository: {args.repo_id}")
+    api.create_repo(repo_id=args.repo_id, repo_type="model",
+                    private=not args.public, exist_ok=False)
+    api.upload_folder(
+        repo_id=args.repo_id,
+        repo_type="model",
+        folder_path=str(Path(__file__).resolve().parent),
+        ignore_patterns=[".git/*", "__pycache__/*", "*.pyc", ".pytest_cache/*", "*.zip"],
+        commit_message="Publish reusable higher-order pruning API",
+    )
+    print(f"Published model repository: https://huggingface.co/{args.repo_id}")
+
+
+if __name__ == "__main__":
+    main()