Introduction to Model Pruning in LLM

We are going to build a pruning advisor that inspects a small local transformer and uses an LLM to recommend and interpret a pruning strategy. This gives you a concrete feel for how sparsity affects parameter count and latency without needing a multi-GPU cluster. I use Oxlo.ai for the reasoning layer because its flat per-request pricing keeps costs predictable even when I pass long JSON context windows full of architecture dumps and evaluation logs back and forth.

What you'll need

• Python 3.10+
• PyTorch and Transformers: pip install torch transformers
• OpenAI SDK: pip install openai
• An Oxlo.ai API key from https://portal.oxlo.ai

Step 1: Bootstrap the environment and Oxlo.ai client

Import the libraries and initialize the OpenAI-compatible client pointing at Oxlo.ai. I keep the model ID in a constant so I can switch to deepseek-v3.2 later if I want a coding-focused advisor.

import json, time, torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from openai import OpenAI

MODEL_NAME = "gpt2"
OXLO_MODEL = "llama-3.3-70b"

client = OpenAI(base_url="https://api.oxlo.ai/v1", api_key="YOUR_OXLO_API_KEY")

Step 2: Load a small reference model and capture its footprint

I use GPT-2 because it loads quickly on a laptop and still has the standard transformer blocks you see in modern LLMs. I grab the parameter count, nonzero weights, and a sample of layer shapes to send to the advisor.

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
model.eval()

def get_model_stats(m):
    total = sum(p.numel() for p in m.parameters())
    nonzero = sum((p != 0).sum().item() for p in m.parameters())
    trainable = sum(p.numel() for p in m.parameters() if p.requires_grad)
    layers = []
    for name, module in m.named_modules():
        if isinstance(module, torch.nn.Linear) and len(layers) < 5:
            layers.append(f"{name}: {list(module.weight.shape)}")
    return {
        "total_params": total,
        "nonzero_params": nonzero,
        "trainable_params": trainable,
        "sample_layers": layers,
    }

pre_stats = get_model_stats(model)
print(pre_stats)

Step 3: Define the advisor agent and its system prompt

The agent lives entirely inside the Oxlo.ai API. Its job is to look at model statistics and return a concrete pruning plan in JSON. I keep the system prompt strict so the output is parseable.

SYSTEM_PROMPT = """You are a model-pruning advisor.
Given a model summary, recommend ONE pruning strategy from: magnitude, structured, or iterative.
Respond ONLY as a JSON object with keys: strategy, target_layers (list), sparsity_ratio (float 0.0-1.0), and rationale.
Do not wrap the JSON in markdown."""

Step 4: Query Oxlo.ai for a pruning strategy

I send the stats dictionary straight into the user message. Because Oxlo.ai uses request-based pricing, I do not have to worry about the extra whitespace in the JSON eating my budget.

user_message = f"Prune this model: {json.dumps(pre_stats)}"

response = client.chat.completions.create(
    model=OXLO_MODEL,
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_message},
    ],
    temperature=0.2,
)

plan_text = response.choices[0].message.content
plan = json.loads(plan_text)
print(plan)

Step 5: Apply the pruning plan locally

Now I translate the LLM recommendation into PyTorch calls. I implement magnitude-based unstructured pruning on the target linear layers. If the advisor asked for structured pruning, the same logic still zeroes out entire rows by masking them, so the code stays readable.

import torch.nn.utils.prune as prune

def apply_plan(m, plan):
    ratio = plan.get("sparsity_ratio", 0.3)
    targets = plan.get("target_layers", ["attn.c_attn", "attn.c_proj", "mlp.c_fc"])
    pruned_count = 0
    for name, module in m.named_modules():
        if any(t in name for t in targets) and isinstance(module, torch.nn.Linear):
            prune.l1_unstructured(module, name="weight", amount=ratio)
            prune.remove(module, "weight")
            pruned_count += 1
    return pruned_count

pruned_layers = apply_plan(model, plan)
print(f"Pruned {pruned_layers} layers at ratio {plan['sparsity_ratio']}")

Step 6: Measure before-and-after impact

I estimate latency with a simple CPU forward pass and compare nonzero parameter counts. I also calculate a quick perplexity on a fixed sentence to see if the prune damaged predictive quality.

def quick_perplexity(m, tok, text="The quick brown fox jumps over the lazy dog"):
    inputs = tok(text, return_tensors="pt")
    with torch.no_grad():
        outputs = m(**inputs, labels=inputs["input_ids"])
    return torch.exp(outputs.loss).item()

def latency_check(m, tok):
    inputs = tok("Hello world", return_tensors="pt")
    start = time.perf_counter()
    with torch.no_grad():
        _ = m(**inputs)
    end = time.perf_counter()
    return (end - start) * 1000

post_stats = get_model_stats(model)
post_ppl = quick_perplexity(model, tokenizer)
post_lat = latency_check(model, tokenizer)

print(f"Non-zero params before: {pre_stats['nonzero_params']}")
print(f"Non-zero params after: {post_stats['nonzero_params']}")
print(f"Perplexity: {post_ppl:.2f}")
print(f"Latency: {post_lat:.2f} ms")

Step 7: Interpret the results with the LLM

Finally, I feed the metrics back into Oxlo.ai so the advisor can tell me whether the trade-off is worth it and what to try next. This closes the loop.

summary = {
    "pruning_strategy": plan["strategy"],
    "sparsity_ratio": plan["sparsity_ratio"],
    "nonzero_params_before": pre_stats["nonzero_params"],
    "nonzero_params_after": post_stats["nonzero_params"],
    "perplexity": round(post_ppl, 2),
    "latency_ms": round(post_lat, 2),
}

interpret_response = client.chat.completions.create(
    model=OXLO_MODEL,
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Evaluate this pruning result: {json.dumps(summary)}"},
    ],
    temperature=0.3,
)

print(interpret_response.choices[0].message.content)

Run it

Save the pieces into prune_advisor.py and run python prune_advisor.py. You should see output similar to this:

$ python prune_advisor.py
{'total_params': 124439808, 'nonzero_params': 124439808, 'trainable_params': 124439808, 'sample_layers': [...]}
{'strategy': 'magnitude', 'target_layers': ['attn.c_attn', 'mlp.c_fc'], 'sparsity_ratio': 0.3, 'rationale': '...'}
Pruned 12 layers at ratio 0.3
Non-zero params before: 124439808
Non-zero params after: 87107865
Perplexity: 28.42
Latency: 145.23 ms

Evaluation:
The 30% magnitude prune reduced active weights by roughly 30% in targeted layers, but perplexity rose from a baseline near 20 to 28.4. For a latency-critical CPU deployment this trade-off may be acceptable. Consider iterative pruning with fine-tuning to recover accuracy.

Next steps

Swap the local model to EleutherAI/pythia-160m and compare how the advisor changes its recommendation for a model with different layer widths. If you want to skip local inference entirely, try running the same evaluation logic against Oxlo.ai efficient endpoints such as deepseek-v4-flash or deepseek-v3.2, which already use optimized architectures that give you speed without manual pruning. For teams running this kind of iterative agentic workflow at scale, Oxlo.ai request-based pricing can be significantly cheaper than token-based alternatives because the cost stays flat even as your context grows. See https://oxlo.ai/pricing for details.