Sheldon Liang, Ali Taghibakhshi, Saurav Muralidharan, Sharath Turuvekere Sreenivas, Ashwath Aithal, Pavlo Molchanov, Nima Tajbakhsh
7 min
Abstract
LLM pruning must preserve quality while meeting hardware- and workload-specific constraints on latency, throughput, and memory. Existing pruning workflows often target parameter counts, requiring practitioners to manually select model sizes that meet latency, throughput, and GPU memory limits. However, parameter count alone does not predict these metrics, which depend on the model architecture, workload, hardware, and serving engine. We introduce AutoCompress, a hardware-guided compression framework that directly incorporates such constraints into differentiable architecture search. AutoCompress jointly optimizes global and per-layer structure to produce heterogeneous architectures, guided by serving-engine cost models validated against end-to-end latency, throughput, and memory measurements. We evaluate AutoCompress on long-context local coding, batched datacenter serving, and latency-sensitive in-game interaction. Compared to the state-of-the-art Minitron pruning framework, AutoCompress improves average accuracy across eight benchmarks by 6.57, 2.35, and 0.57 percentage points on RTX 5070, RTX 5080, and H100, respectively. AutoCompress also compresses Nemotron 3.5 Super from 120B to 74B parameters for single-B200 BF16 serving at 128K context, retaining 95% of its mean benchmark score.
Sam: Measuring every complete candidate sounds expensive. How do they make the search manageable?
Alex: They profile short models containing different layer configurations inside the actual serving engine. Those measurements populate a timing lookup. A candidate’s estimated runtime adds its layer costs and the overhead outside the layers. Memory estimates include weights, stored request state, temporary workspaces, and runtime buffers.
Sam: That last part matters: fitting the weights is not the same as fitting the running model. But they still need to estimate which cuts damage quality.
Alex: They rank components by activation-based importance and reorder the parent model’s weights. More important components come first. Then they construct a supernet: shared weights supporting many candidate architectures. Smaller candidates use subsets of those weights rather than requiring separate training from scratch.
Sam: Where does the decision-making happen? Sharing weights gives you candidates, but not a reason to prefer one.
Alex: An architecture router learns which dimensions to retain. During search, its choices are softened, so training can gradually adjust them. The parent teaches the candidates through knowledge distillation: training smaller models to imitate its output probabilities. Resource penalties push choices toward the deployment limits.
Sam: Does that reward being as small as possible? That could leave usable memory idle while throwing away quality.
Alex: The penalty discourages both exceeding the budget and leaving too much unused. The goal is to retain capacity near the limit, not minimize size. Once choices settle, the system extracts the selected architecture from the reordered parent weights and distills it further.
Sam: From the parent weights, not the weights trained during the search?
Alex: Yes. The search-trained shared weights help evaluate architectural choices. Recovery starts from the corresponding retained parent weights. That distinction separates finding a structure from training the final standalone model.
Sam: Give me an example where hardware costs change what the system decides to preserve.
Alex: The models use a mixture of experts: specialist blocks selected for each token. They also have a shared expert block. In the local-coding search, the shared block stays full-sized. Shrinking it saves relatively little serving time and memory compared with similarly shrinking the routed expert blocks.
Sam: So the search can preserve something useful because removing it barely helps the deployment budget. That is more specific than cutting every layer equally. How strong is the comparison?
Alex: The authors strengthen Minitron into a baseline called Minitron-plus-plus. It gets the same serving-cost models, deployment budgets, and importance-sorted parent. Both methods receive matched recovery training. But the baseline keeps uniform layer dimensions and retains all layers.
Sam: Then the comparison bundles a different search procedure with a more flexible architecture space. It cannot tell us that the search algorithm alone causes the gain.
Alex: That is the right limit on the interpretation. The evidence supports the complete AutoCompress workflow against this deployment-adapted baseline. Across eight benchmarks, the average advantage is about six-point-six percentage points for the gaming target. The datacenter advantage is only about half a point.
Sam: So the benefit is not equally large everywhere. Did the models actually satisfy their serving requirements, rather than just satisfy the predictor?
Alex: The selected architectures meet their deployment budgets in hardware measurements. The consumer models satisfy both response-time and memory limits; the datacenter model reaches its throughput target. AutoCompress also uses fewer architecture-selection tokens, while recovery training is matched within each target.
Sam: What about a much larger model? Does the paper test whether this is more than a consumer-model recipe?
Alex: A separate case study reduces Nemotron Super from roughly one hundred twenty billion parameters to seventy-four billion. The compressed model supports long-context serving on a single B200 graphics processor. It retains about ninety-five percent of the parent’s mean benchmark score.
Sam: Is that another controlled win against the baseline, or a demonstration that the larger model can be made to fit?
Alex: The latter. They do not run Minitron-plus-plus there because of the training cost. And transfer is another boundary: the cost models are calibrated to particular deployment settings. The authors acknowledge that architecture changes can alter kernel selection, cache behavior, and overhead, limiting additive timing estimates.
Sam: Who should read the full document, and where would you send them first?
Alex: Researchers building compression systems or deploying constrained language models. Start with Accuracy Under Deployment Budgets, alongside the Minitron-plus-plus Baseline section, to understand the comparison. Then read the latency and memory cost models. Those details determine whether this approach matches your serving environment.
Sam: And for someone who just wants the lesson to carry into their next deployment discussion?
Alex: Compress for the machine and workload, not just the parameter count.
Sam: Keep the deployment budget in the conversation.