NASPipe: high performance and reproducible pipeline parallel supernet training via causal synchronous parallelism

Searching for a neural network architecture means training thousands of overlapping subnetworks in a strict order — parallelise that across GPUs and they start reading each other's half-written layers, so the search stops being repeatable.

Why this problem is worth solving

Making it faster still means running several of those paths at once, which turns the problem into one of

What people did before, and where it broke

• Train each candidate architecture separately. The original approach, and the reason architecture search was once out of reach for most people.
• Give each subnetwork its own GPU. The natural parallelisation, used by existing supernet frameworks with a central parameter server for synchronisation. Two problems: the server does not scale when subnetworks synchronise frequently, and modern supernets are large enough that a single subnetwork already exceeds one GPU.
• Use bulk synchronous parallelism. Insert a barrier after every batch of parallel work, as pipeline systems built for single large models do. It preserves an ordering, but not the one that matters here, and the barriers cost a great deal of parallelism.
• Use asynchronous or stale synchronous parallelism. Let workers proceed without waiting, as one prominent pipeline system does. Very efficient — its pipeline sits idle only ten percent of the time — and it makes no attempt to preserve any dependency between subnetworks at all.
• Swap the model between CPU and GPU memory. Necessary because supernets exceed GPU capacity. Existing systems that do this assume the sequence of operations is known in advance so they can prefetch. In supernet training it is generated at runtime.
• The gap. The synchronisation models on offer were designed for parallelising one model’s own gradient computation, not for ordering constraints between different models that happen to share layers. Neither serialising everything nor ignoring the constraint is acceptable.

How NASPipe preserves the order and still goes fast

The idea

Say precisely what parallel execution is allowed to break. The contribution is a synchronisation model, which the authors call causal synchronous parallelism, and its value is that it states the constraint exactly rather than approximating it from either side. The rule is: if an earlier subnetwork and a later one both select the same layer, the later one may not read or write that layer until the earlier one’s write has finished. Everything else is free — workers compute locally and asynchronously and overlap with communication. That is narrower than a barrier, which orders everything whether or not it needs ordering, and stricter than asynchronous updates, which order nothing. Naming a correctness property that exists only in the relationship between concurrent actors, so that a system can enforce exactly it, is the same discipline behind defining the cross-thread inconsistencies that arise when concurrency meets persistence, covered elsewhere on this site.

A bigger search space is easier to parallelise, not harder. The performance insight inverts the usual intuition. Two subnetworks conflict when they happen to select the same layer, and in a large supernet with many choice blocks and many candidates per block, two subnetworks drawn close together in the sequence are unlikely to overlap much. So the scheduler does not process the sequence in order and stall; it reaches forward and runs whatever has no unfinished precedent. The consequence shows up directly in the results: against one baseline the improvement is a mere ten percent on the smallest search space and 6.8 times on the largest. The choice of pipeline parallelism rather than one-subnetwork-per-GPU matters here too, for two reasons — dependency checking and synchronisation become local decisions on each GPU rather than round trips to a central server, and modern supernets have subnetworks too large to fit on one device anyway. Reconsidering how a pipeline should be partitioned and scheduled once the workload stops looking like a single fixed model is the same exercise performed by work adapting pipeline parallelism to recomputation and partitioning trade-offs, covered elsewhere on this site.

Prefetching requires knowing the future, and here the future is computed. The supernet does not fit in GPU memory, so it lives in CPU memory and layers are swapped in and out. Existing systems that do this rely on knowing the operator sequence in advance — but here the sequence comes from the exploration algorithm at runtime, and the scheduler deliberately reorders it. The way out is an observation about hardware rather than about the algorithm: operator execution times on GPUs are close to deterministic, so from the state of the current stage plus information passed along the pipeline, the next few scheduled tasks can be forecast. The predictor even runs the scheduler speculatively — before a backward pass it adds that pass to the finished set and re-runs the scheduling decision, since completing it will unblock something. It is right about ninety percent of the time, which is enough to avoid most forced stalls. Leaning on the fact that deep learning operator timings are predictable enough to plan against is the same foundation used by work forecasting GPU performance for training and inference, covered elsewhere on this site.

Following two conflicting subnetworks

The forward pass and the backward pass of a stage are separate scheduling units. Each is identified by whether it is forward or backward, which subnetwork it belongs to, and which stage it is.

Why it works

• The constraint is narrower than a barrier. Only subnetworks that actually share a layer must be ordered → everything else runs concurrently → so correctness is preserved without paying the full serialisation a bulk barrier imposes.
• Conflicts thin out as the space grows. More choice blocks and more candidates per block make overlap between nearby subnetworks less likely → the scheduler can reach further ahead → which is why the advantage over the barrier-based baseline goes from ten percent on the smallest space to 6.8 times on the largest.
• Finishing a write releases many readers. A backward pass completes a subnetwork’s update to its layers → prioritising backward tasks removes precedence constraints from everything downstream → enlarging the scheduling search space rather than merely draining the current one.
• Pipelining makes the checking local. Each GPU owns a stage and checks dependencies for it → no central synchronisation server sits on the path → which matters because subnetworks synchronise frequently and a central server does not scale.
• Deterministic operator timing makes the schedule forecastable. The next tasks can be predicted from pipeline state at around 90% accuracy → layers are prefetched before they are requested → so keeping the supernet in CPU memory costs little, and larger batches fit on the GPU.

What it costs (trade-offs and scope)

• It is not faster than every baseline. Against one recent pipeline system the range is 0.77 to 1.5 times — meaning NASPipe is slower on some workloads. The 7.8-fold headline is against a different, barrier-based baseline.
• Enforcing the dependency costs pipeline efficiency. NASPipe’s pipeline sits idle 39% to 43% of the time. The asynchronous baseline achieves 10%, precisely because it enforces nothing. That gap is the price of the correctness property.
• The baselines are not doing the same job. The paper states plainly that the comparison systems do not enforce causal dependencies. NASPipe is therefore doing strictly more work and still winning on throughput in most cases — but the comparison is between a correct system and incorrect ones, which is worth holding in mind when reading a speedup number.
• Two of the six search spaces were shrunk to let the baselines run. For those, the authors report that their context management gained little advantage, and go further, noting that on such small search spaces the existing data-parallel approach would likely be comparable to everything here.
• Prediction misses one time in ten. When the forecast is wrong, execution stalls for a synchronous layer transfer from CPU memory.
• Scaling is sub-linear and degrades. From 4 to 16 GPUs, communication grows with pipeline length and the idle fraction rises from 39% to 42%. The baselines scale worse, but nothing here scales well.
• It requires an imperative training framework. Systems with static execution graphs cannot integrate with it, though the authors note the popular frameworks all support imperative execution. The optimisations also target the currently dominant search paradigms rather than the full variety of methods.

The evaluation

Check yourself

Sources

The paper is “NASPipe: High Performance and Reproducible Pipeline Parallel Supernet Training via Causal Synchronous Parallelism” by Shixiong Zhao, Fanxin Li, Xusheng Chen, Tianxiang Shen, Li Chen, Sen Wang, Nicholas Zhang, Cheng Li and Heming Cui, published at ASPLOS 2022, pages 374-387, DOI 10.1145/3503222.3507735. The authors are at the University of Hong Kong, Huawei Technologies and the University of Science and Technology of China.

This explainer is written from the full text, read from the author-hosted copy at the University of Hong Kong. All figures above come from the paper’s own sections 1 through 5: the roughly 500-fold reduction in search cost from the supernet paradigm, the definition of causal synchronous parallelism with its dependency preservation and concurrent computation properties, the scheduler’s backward-first priority and its sub-ten-millisecond decision cost, the context predictor’s roughly 90% accuracy, the throughput ranges of 1.1 to 7.8 times against the barrier-based system, 0.87 to 6.5 times against the asynchronous system and 0.77 to 1.5 times against the memory-swapping system, the pipeline idle fractions of 39% to 43% against 10% for the asynchronous baseline, the growth of advantage with search space size from 1.1 to 6.8 times, the reproducibility results across three runs on 4 and 8 GPUs, and the scaling behaviour from 4 to 16 GPUs. The authors state their own limitations in section 5.5, including the requirement for an imperative training framework, and note in section 5.1 that two search spaces were reduced in size so that baseline systems could run at all.