Hyperscale Hardware Optimized Neural Architecture Search

Automated search for neural network architectures works well in research papers and almost never survives the trip into a production datacentre — the data is different, the constraints are different, and the winning model has to be retrained and hand-tuned before anyone will ship it; this paper from Google rebuilds the whole pipeline for that setting, searching directly on live production traffic that never touches a disk, scoring candidates with a learned model of how fast the hardware will actually run them, and reports 22 to 56 percent better performance and 17 to 25 percent better energy efficiency at neutral or better model quality, deployed at zettaflop scale.

Why this problem is worth solving

The method being adapted has a specific and slightly odd shape.

The hardware property that decides the winners is not the one people optimise for.

What people did before, and where it broke

• Design models by hand and tune them afterwards. The default. Experts pick an architecture for accuracy, then engineers optimise the implementation. Where it broke. The interactions this paper cares about run the other way — the architecture itself decides the balance between compute, memory and network, and rebalancing that is described as very hard for humans. On the production recommendation model, a ten percent gain the search found in days is said to match what a team of ten or more experts had historically achieved over months.
• Multi-trial search. Sample an architecture, train it properly, measure it, repeat. Straightforward and accurate. Where it broke. Cost. Each trial is a full training run, and the models here train on 128 accelerator chips. It is not affordable at this scale.
• Use operation counts as the speed objective. Cheap, hardware-independent, easy to differentiate. Where it broke. The correlation to real performance is poor — over 400 percent error in prior work — because operation counts ignore operational intensity entirely.
• Measure candidates on real hardware during the search. The most accurate signal, and standard practice for multi-trial search. Where it broke. In one-shot search the individual candidate does not physically exist; it is a slice of a super-network. And a reward is needed every step, on a timescale of tens to hundreds of milliseconds. Simulators are too slow for the same reason.
• Train a learned performance model offline. Prior work does this and it works at research scale. Where it broke. The training data. Existing approaches need thousands of hardware measurements for search spaces orders of magnitude smaller than this one; scaled up it would take millions, on models that occupy hundreds of expensive accelerators each.
• The gap. Every piece of the research pipeline assumes something production does not provide — a stored, splittable dataset; the freedom to retrain; a cheap way to evaluate candidates; and a single speed objective rather than a launch checklist. The question nobody had answered is what architecture search looks like when it is rebuilt around those constraints instead of excused from them.

How a search runs on live production traffic

The idea

The two-dataset requirement is replaced by an ordering guarantee. Every prior search method needs a training set and a validation set, because learning the architecture is a form of evaluation and evaluating on data the weights were trained on causes overfitting. Splitting a live production stream into two statistically stable halves inside an in-memory pipeline is prohibitively expensive, so the paper drops the split and keeps the property it was protecting. With enough traffic, every sample can be used exactly once, and the pipeline enforces that the architecture decision consumes each batch before the weight update does. The architecture is therefore always judged on data the shared weights have never seen — which is what a validation set is for — and the two kinds of learning collapse into one step. Getting parallel work to respect an ordering dependency between weight-sharing sub-networks, rather than serialising or silently violating it, is the same knot the pipeline-parallel supernet trainer covered elsewhere on this site had to untie from the systems side.

The search space is where the hardware knowledge lives. This is the part that makes it an architecture paper rather than a machine learning one. Two convolution block variants illustrate it: the unfused one performs less total arithmetic, and the fused one performs more but with higher operational intensity and therefore a higher rate. Which is faster depends on the channel depth — at 32 channels the fused block wins on latency, at 128 it loses despite still having the better intensity. Neither block is correct in general, so both go in the space and the search picks per layer. The same logic shapes the recommendation space, where embedding layers are memory- and network-bound while the dense layers are compute-bound, and the balance between them sets both the hardware utilisation and the model’s balance between memorisation and generalisation. Searching a space defined by what the hardware is good at, rather than by what looks reasonable on paper, is the same instinct as the full-stack accelerator search covered elsewhere on this site, which searched the datapath, the schedule and the compiler together.

A learned performance predictor becomes affordable by being pre-trained on simulations. The inner loop needs a speed number every step, in tens to hundreds of milliseconds, for a candidate that does not physically exist. Operation counts are too inaccurate, hardware measurement is impossible, and training a predictor on real measurements alone would need millions of runs on hundreds of chips. The answer is to split the training into two phases with different currencies — a million cheap simulated samples to learn the shape of the non-linear performance surface, then twenty expensive real measurements to correct its calibration. The reported error drops from between 14.7 and 42.9 percent after pre-training alone to between 1.05 and 3.08 percent after fine-tuning. Learning to predict hardware performance instead of measuring it, and structuring the predictor so the learned part has as little to do as possible, is the same strategy the work on forecasting performance for unavailable accelerators covered elsewhere on this site pursues in a setting where the hardware cannot be touched at all.

Walking through one search step

This ordering is the entire justification for having no validation set. Without it, the architecture would be scored on data the weights had already fitted, and the search would overfit and generalise poorly.

Why it works

• Live traffic removes a whole stage of the pipeline. The usual reason a searched model needs retraining is that it was found on a different data distribution from the one it will serve → searching on production traffic directly means the distribution gap does not exist → so the model that comes out of the search is closer to deployable, and the manual intervention that normally erodes its quality is avoided.
• The ordering rule buys generalisation without a second dataset. A validation split exists to keep architecture decisions honest → the property that actually matters is that the architecture is judged on data the weights have not fitted → sequencing each batch so the architecture sees it first delivers exactly that property, and at hyperscale there is enough traffic that no sample ever needs reusing.
• One-shot plus reinforcement learning is the only combination that fits. Multi-trial search costs a full training run per candidate → gradient-based one-shot search must compute gradients for every sub-network on every step rather than only the sampled one → reinforcement learning over a shared super-network activates one candidate per core per step, which is what makes hundreds of parallel samples affordable.
• Two-phase training makes the performance model cheap and accurate at once. Simulation is inexpensive and captures the non-linear shape of the performance surface but drifts from reality → real measurements are accurate but cost hundreds of chips each → pre-training on a million simulated samples and fine-tuning on twenty real ones takes the calibration from the simulator and the accuracy from the hardware.
• A one-sided penalty matters precisely because there are several targets. With a single performance objective, penalising deviation in either direction is harmless → with several, candidates meeting all targets at once are rare, so the search needs the ones that overshoot on some target to move through the space → not punishing them is what keeps the sparse region navigable, worth up to 13 percent better step time at comparable quality.

What it costs (trade-offs and scope)

• The single-step algorithm does not scale down. The paper states plainly that two-step learning is still needed for small research datasets, because using each sample once requires having far more data than the model needs. The central algorithmic contribution is contingent on hyperscale, not portable from it.
• The headline range mixes very different comparisons. The 22 to 56 percent spans a 1.54 times speedup on a research vision transformer family and a 1.22 times gain on production recommendation models. The convolutional family gain is 5 to 6 percent averaged across eight models, of which five were left completely unchanged by the search; only the largest three moved, by 14 to 17 percent.
• The ablation shows one change carrying almost all of the vision speedup. By the paper’s own table, going from the baseline to the searched model raises throughput from 101 to 186 images per second, and every bit of that comes from shrinking pre-training resolution from 224 to 160 pixels — a change that alone costs 1.4 percent top-1 accuracy, which the extra convolution depth and the changed activation function then buy back. Neutral quality at 1.84 times the speed on the largest model is real, but it decomposes into a resolution cut and two compensations.
• The performance model is pre-trained on a simulator’s opinion. The simulator is in-house, does not model host processors, and assumes models run entirely on accelerators. Pre-training alone gives 14.7 to 42.9 percent error against production measurements, so the twenty fine-tuning samples are doing the decisive work.
• A new performance model is needed whenever the hardware or the search space changes, and reusing one pre-trained model across domains is said to cost significant accuracy. A universal pre-trained model is named as future work.
• Nothing here can be reproduced outside Google. The traffic, the recommendation baseline, the simulator and the accelerators are all internal. Two model families were open-sourced, which is the reproducible part.
• The reward improvement is modest and conditional. Up to 13 percent better training step time, up to 0.4 percent better quality, 1.6 percent smaller serving memory — and no difference at all when there is only one performance objective.
• Quality is protected but not always improved for free. Some production models were deliberately allowed to lose performance in order to gain quality, which is the right call and worth noticing when reading the averages.
• The search costs about 2.5 times a normal training run. That is cheap only because it happens once and the resulting model is served for a long time afterwards — reported as under 0.03 percent of the accelerator hours the downstream serving and research jobs consume.

The evaluation

Check yourself

Sources

The paper is “Hyperscale Hardware Optimized Neural Architecture Search” by Sheng Li, Garrett Andersen, Tao Chen, Liqun Cheng, Julian Grady, Da Huang, Quoc V. Le, Andrew Li, Xin Li, Yang Li, Chen Liang, Yifeng Lu, Yun Ni, Ruoming Pang, Mingxing Tan, Martin Wicke, Gang Wu, Shengqi Zhu, Parthasarathy Ranganathan and Norman P. Jouppi, published at ASPLOS 2023, pages 343-358, DOI 10.1145/3582016.3582049. It is open access under a Creative Commons licence, and two model families are open-sourced.

This explainer is written from the full 16-page text. Everything above comes from the paper’s own sections: the taxonomy of search strategies, algorithms, spaces and objectives, and the three production challenges of deployment, efficiency and scale; the in-memory pipeline motivated by data privacy rules, the argument for why every prior method needs two datasets, and the single-step algorithm with its ordering guarantee, its three parallel stages and its explicit limitation to large datasets; the search space design, including the first recommendation-model space and super-network for this style of search, the hybrid of fine-grained and coarse-grained weight sharing across embedding widths, vocabulary sizes and dense layers, the space size, and the fused and unfused convolution block comparison with its roofline and latency measurements at 32 and 128 channels; the single-sided reward function, its comparison against the prior absolute-value form, and the measured gains in step time, quality and serving memory; the two-phase performance model with its architecture, its million pre-training samples, its twenty fine-tuning measurements, the error figures before and after fine-tuning, and the description of the in-house simulator and what it does not model; and the evaluation, covering the vision transformer results with their ablation table and hardware analysis, the recommendation model load-imbalance finding, the convolutional family geometric means with the note that the smaller models were unchanged, the power and energy analysis including why the convolutional family shows no power gain, the production deployment figures for vision and recommendation models, and the accounting of what the search itself costs.