LLM Model Pruning: Techniques and Best Practices

Model pruning cuts inference costs by removing redundant weights, but hand-picking which layers to drop is tedious and error-prone. I built a lightweight pruning workbench that uses Oxlo.ai to recommend a pruning strategy and judge output quality, while the actual weight surgery happens locally on a small open model. Because Oxlo.ai uses flat per-request pricing, running a dozen evaluator calls with full context does not inflate the bill the way token-based providers do.

What you'll need

• Python 3.10+
• pip install openai torch transformers datasets
• An Oxlo.ai API key from https://portal.oxlo.ai
• About 1 GB of free disk space for the 135 M parameter test model

Step 1: Set up the Oxlo.ai client and local environment

I start by importing the libraries and initializing the OpenAI-compatible client pointing at Oxlo.ai. I also define the local test model and the sparsity target we will apply.

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

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

MODEL_ID = "HuggingFaceTB/SmolLM-135M"
TARGET_SPARSITY = 0.30

Step 2: Inspect the model and generate a pruning plan

I load the 135 M parameter model and extract its layer configuration. Instead of guessing which layers to prune, I send the architecture summary to Llama 3.3 70B on Oxlo.ai and ask for a structured plan. Here is the system prompt I use for the pruning strategist.

SYSTEM_PROMPT = """You are a senior ML engineer specializing in model compression. The user will provide a transformer architecture summary and a target sparsity ratio. Respond with a JSON object containing a single key \"prune_layers\" whose value is a list of integer layer indices to prune. Prefer pruning later MLP layers over early attention layers. Respond with only the JSON object."""

Next, I build the summary and call Oxlo.ai.

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32)

num_layers = len(model.model.layers)
hidden_size = model.config.hidden_size
intermediate_size = model.config.intermediate_size

arch_summary = (
    f"Model: {MODEL_ID}\n"
    f"Layers: {num_layers}\n"
    f"Hidden size: {hidden_size}\n"
    f"Intermediate size: {intermediate_size}\n"
    f"Target sparsity: {int(TARGET_SPARSITY * 100)}%"
)

response = client.chat.completions.create(
    model="llama-3.3-70b",
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": arch_summary},
    ],
)

plan = json.loads(response.choices[0].message.content)
layers_to_prune = plan["prune_layers"]
print("Layers selected for pruning:", layers_to_prune)

Step 3: Apply structured pruning to the local model

I clone the original model so I can compare later. For each layer index the agent selected, I compute the L2 norm of every neuron in gate_proj and zero out the weakest 30 percent. I propagate the same mask to up_proj and down_proj to keep dimensions consistent.

import copy

pruned_model = copy.deepcopy(model)

def prune_mlp_layers(target_model, layer_indices, sparsity_ratio):
    layers = target_model.model.layers
    for idx in layer_indices:
        if idx < 0 or idx >= len(layers):
            continue
        mlp = layers[idx].mlp
        gate = mlp.gate_proj.weight.data
        norms = torch.norm(gate, dim=1)
        k = int(sparsity_ratio * len(norms))
        if k == 0:
            continue
        threshold = torch.kthvalue(norms, k).values
        mask = norms > threshold
        gate[~mask] = 0
        mlp.up_proj.weight.data[~mask] = 0
        mlp.down_proj.weight.data[:, ~mask] = 0
    return target_model

pruned_model = prune_mlp_layers(pruned_model, layers_to_prune, TARGET_SPARSITY)
print(f"Pruned {len(layers_to_prune)} layers at {TARGET_SPARSITY:.0%} sparsity.")

Step 4: Evaluate outputs with Oxlo.ai as a judge

I generate answers from both the original and pruned models on a small validation set. Then I send each pair to Oxlo.ai for a side-by-side quality score. Because Oxlo.ai pricing is per request, I can include the full prompt and both outputs without worrying about token count.

test_prompts = [
    "Explain model pruning in one sentence.",
    "What is the capital of France?",
    "Write a Python function to reverse a list.",
]

def generate_answer(target_model, prompt, max_new_tokens=60):
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = target_model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

original_outputs = [generate_answer(model, p) for p in test_prompts]
pruned_outputs = [generate_answer(pruned_model, p) for p in test_prompts]

results = []
for prompt, orig, pruned in zip(test_prompts, original_outputs, pruned_outputs):
    judge_input = (
        f"User prompt: {prompt}\n\n"
        f"Original output: {orig}\n\n"
        f"Pruned output: {pruned}\n\n"
        "Score the pruned output on coherence and factual consistency relative to the original. "
        "Respond with JSON: {\"score\": int, \"reason\": str}"
    )

    response = client.chat.completions.create(
        model="llama-3.3-70b",
        messages=[
            {"role": "system", "content": "You are an expert evaluator of language model outputs."},
            {"role": "user", "content": judge_input},
        ],
    )

    verdict = json.loads(response.choices[0].message.content)
    results.append({"prompt": prompt, "score": verdict["score"], "reason": verdict["reason"]})
    print(f"Prompt: {prompt[:40]}... Score: {verdict['score']}/10")

Step 5: Aggregate scores and decide whether to keep the prune

I average the coherence scores and print a recommendation. If the average is above 8, I save the pruned weights. Otherwise, I tighten the sparsity target and rerun.

avg_score = sum(r["score"] for r in results) / len(results)
print(f"\nAverage quality score: {avg_score:.1f}/10")

if avg_score >= 8.0:
    print("Pruning is viable. Save the weights with torch.save(pruned_model.state_dict(), 'pruned_model.pt').")
else:
    print("Quality dropped too much. Reduce the sparsity ratio or prune fewer layers.")

Run it

Copy the blocks above into a single file named prune_agent.py, replace YOUR_OXLO_API_KEY, and run python prune_agent.py. Here is what the output looks like on my machine.

$ python prune_agent.py
Layers selected for pruning: [24, 26, 27, 28, 29]
Pruned 5 layers at 30% sparsity.
Prompt: Explain model pruning in one sentence... Score: 9/10
Prompt: What is the capital of France?... Score: 10/10
Prompt: Write a Python function to reverse a list... Score: 8/10

Average quality score: 9.0/10
Pruning is viable. Save the weights with torch.save(pruned_model.state_dict(), 'pruned_model.pt').

Wrap-up and next steps

This agent gives me a repeatable way to test pruning hypotheses without burning tokens on long evaluation prompts. Two concrete next steps: wire the judge loop into a CI test that fails if the score drops below a threshold, and try the same workflow on larger checkpoints by offloading reference generation to deepseek-v3.2 on Oxlo.ai while your laptop handles the sparse forward pass.