AutoComp: Agent-Based Structured Pruning for Large Language Models

Installation

conda env create -f environment.yml
conda activate autocomp

Set your API key before running search:

export OPENAI_API_KEY="your_openai_api_key"

If you use gated Hugging Face models, log in first:

huggingface-cli login

Database-building scripts require CUDA GPUs.

Smoke Test

Run a lightweight check without GPUs or API calls:

python tests/smoke_test.py

Optionally test local Hugging Face model loading and one forward pass:

AUTOCOMP_SMOKE_MODEL="Qwen/Qwen2.5-3B-Instruct" python tests/smoke_test.py

Run

Each supported model has two scripts:

1. Build the pruning database.
2. Run search and benchmark the top 3 policies.

Llama-3.1-8B

bash scripts/build_db_llama31_8b.sh
bash scripts/search_benchmark_llama31_8b.sh

Qwen2.5-3B-Instruct

bash scripts/build_db_qwen25_3b.sh
bash scripts/search_benchmark_qwen25_3b.sh

Qwen2.5-7B-Instruct

bash scripts/build_db_qwen25_7b.sh
bash scripts/search_benchmark_qwen25_7b.sh

Mistral-Nemo-Instruct-2407

bash scripts/build_db_mistral_nemo_12b.sh
bash scripts/search_benchmark_mistral_nemo_12b.sh

OLMoE-1B-7B-0924

bash scripts/build_db_olmoe_1b7b.sh
bash scripts/search_benchmark_olmoe_1b7b.sh

Common Options

Override defaults with environment variables:

MODEL="meta-llama/Llama-3.1-8B" \
DB_PATH="db_output/llama31_8b/Llama-3.1-8B" \
RUN_ROOT="outputs/llama31_8b" \
TARGET_FLOPS_REDUCTION=0.5 \
AGENT_MODEL="gpt-5-mini" \
AGENT_STEPS=40 \
bash scripts/search_benchmark_llama31_8b.sh

Useful variables:

• MODEL: model name or local model path.
• DB_PATH: pruning database path.
• RUN_ROOT: output directory.
• TARGET_FLOPS_REDUCTION: target compression ratio.
• AGENT_MODEL: OpenAI model for policy search.
• AGENT_STEPS: number of search rounds.
• CUDA_VISIBLE_DEVICES: GPUs to use.
• NPROC_PER_NODE: number of processes for database building.
• PYTHON_BIN: Python executable used by scripts.
• TORCHRUN_BIN: torchrun executable used by database-building scripts.
• ALLOW_CPU_BUILD: set to 1 only for tiny CPU debugging runs with NPROC_PER_NODE=1.

Outputs

Search scripts write:

RUN_ROOT/search_result.json
RUN_ROOT/best_policy.json
RUN_ROOT/debug_trace.json
RUN_ROOT/policies/top1_policy.json
RUN_ROOT/policies/top2_policy.json
RUN_ROOT/policies/top3_policy.json
RUN_ROOT/benchmarks/top1_benchmark_results.json
RUN_ROOT/benchmarks/top2_benchmark_results.json
RUN_ROOT/benchmarks/top3_benchmark_results.json