ProteuS: Simulator for Distributed DNN Training

• ProteuS is a standalone simulator for distributed DNN training that models complex strategies using a hierarchical Strategy Tree for accurate performance prediction.
• It compiles the training strategy into a distributed execution graph, explicitly incorporating communication, scheduling, and recomputation to mirror real-world execution.
• Evaluations show an average prediction error of around 3%, preserving throughput ordering and outperforming prior approaches in modeling compute–communication overlap and bandwidth sharing.

Proteus is a standalone simulator for predicting the performance of distributed DNN training under complex parallelization strategies. It was introduced to address a gap between modern training practice and prior performance models: large DNNs are routinely trained with nested combinations of data parallelism, model parallelism, pipeline parallelism, recomputation, and memory partitioning, while runtime performance depends not only on per-operator costs but also on dynamic interactions such as compute–communication overlap and bandwidth sharing. Proteus addresses this by representing training strategies with a hierarchical “Strategy Tree,” compiling that representation into a distributed execution graph, and simulating runtime execution with a Hierarchical Topo-Aware Executor (HTAE). Across six representative models and three hardware configurations, it reports a $3.0\%$ average prediction error, preserves throughput ordering across candidate strategies, and reduces prediction error by up to $133.8\%$ relative to prior approaches (Duan et al., 2023).

1. Problem formulation and design objectives

Proteus was developed in the context of rapidly growing DNN models whose computation and memory demands necessitate large clusters and increasingly elaborate parallelization strategies. The central modeling problem is that distributed training throughput can no longer be approximated reliably by summing static per-layer compute and communication costs. The paper argues that the underlying assumption of many prior works—namely, that the cost of a single operator only depends on its input and output tensor shape—breaks down once sophisticated combinations of operator-level and subgraph-level parallelism are used. In such settings, communication can overlap with computation, and multiple communication groups can contend for shared physical links, so effective costs depend on global execution context rather than on isolated operators alone (Duan et al., 2023).

This motivates a simulation-based formulation rather than a purely analytical one. Proteus is explicitly presented as the first standalone simulator for complex parallelization strategies. Its purpose is not to search the entire parallelization space, but to evaluate specified strategies accurately enough to compare them, identify bottlenecks, estimate resource requirements, and support automated parallelization systems. A plausible implication is that Proteus occupies a middle layer between low-level cost estimation and full training execution: it is more expressive than static analytical models, but still substantially cheaper than empirical trial runs.

2. Strategy Tree: a unified representation of parallelization

The main abstraction underlying Proteus is the Strategy Tree, a hierarchical representation intended to unify operator-level and subgraph-level parallelism. The tree mirrors the nested structure of a DNN. Leaf nodes correspond to layers, operators, or tensors, while non-leaf nodes correspond to subgraphs composed of multiple layers. Each node stores a parallel configuration. For leaf nodes, that configuration includes a computation config and a memory config; for non-leaf nodes, it includes a schedule config. The computation or memory config contains a partition $\mathcal{P}$ that specifies the degree of parallelism in each dimension, together with a map from partitions to devices. The schedule config includes three parameters: n_micro_batch, max_ongoing_micro_batch, and recomputation (Duan et al., 2023).

This representation is designed to express nested strategies that mix fine-grained operator sharding with coarser subgraph scheduling. The paper distinguishes operator-level strategies, which include computation parallelization and memory optimization, from subgraph-level strategies, which include pipeline parallelism and recomputation. In this formulation, memory optimization is itself handled by partitioning tensor dimensions, enabling representation of ZeRO-style partitioning or activation partitioning within the same abstraction. This suggests that Proteus treats memory layout and execution scheduling as first-class components of the training strategy rather than as secondary implementation details.

The paper illustrates operator parallelism with a linear operator: $$output(b, s, o) = \sum_{h} input(b, s, h) \times weight(o, h).$$ The parallelizable dimensions are $b$ (batch), $s$ (sequence), $o$ (output channel), and $h$ (hidden or reduction). Data parallelism corresponds to splitting the batch dimension; model parallelism corresponds to splitting output or reduction dimensions; hybrid schemes combine these; and “op shard” generalizes them to arbitrary tensor partitioning. A key detail is that Proteus defines strategies explicitly at the tensor level. If the implicit tensor strategy induced by one operator conflicts with the explicit strategy required by another, Proteus accounts for the communication required to transform between them. That choice is central to modeling realistic distributed graphs, since redistribution is often a dominant hidden cost in nontrivial parallelization plans.

3. Compilation into a distributed execution graph

Proteus compiles the Strategy Tree into a distributed execution graph through an execution graph compiler. The first compilation stage partitions the DNN into disjoint subgraphs according to device groups, denoted DevGroups. The compiler traverses the tree breadth-first from the root, and a node can be divided only if its children share some devices. Each resulting subgraph is then compiled into forward and backward execution graphs. Tensors and operators are split into partitions so that each partition resides on, and is executed by, one device (Duan et al., 2023).

When producer and consumer partitions disagree, the compiler automatically inserts communication operators via strategy transformation. The paper states that Proteus uses pattern matching to infer collective communication primitives such as All-Reduce, and falls back to point-to-point communication when needed. This mechanism is important because a large fraction of distributed training overhead arises not from nominal operator execution, but from communication introduced by mismatched partition semantics across adjacent graph regions.

The compiler also introduces control dependencies to enforce the user-specified schedule. Forward subgraphs are made dependent on their corresponding backward subgraphs to control memory usage, and recomputation subgraphs are constrained to execute immediately before backward execution. For recomputation, the compiler creates a structure with two forward and one backward execution subgraph so that one forward copy can be released and later recomputed. This construction makes recomputation explicit in the execution graph rather than treating it as a post hoc correction to memory estimates.

4. HTAE and runtime-behavior simulation

Proteus’s simulation engine is the Hierarchical Topo-Aware Executor. HTAE is organized into two levels. At the top level is a scheduler that interleaves multiple subgraphs according to data and control dependencies while balancing micro-batch parallelism and memory. At the lower level, each subgraph has an executor with three queues: computation, feature communication, and gradient communication. This separation enables simulation of compute–communication overlap while distinguishing feature communication from gradient communication so that one does not unnecessarily block the other. The executor alternates between computation and communication, dispatching dependency-free operations as they become available. HTAE also tracks tensor reads, writes, and reference counts, releasing memory when references drop to zero; this enables prediction of out-of-memory behavior in addition to throughput (Duan et al., 2023).

Two runtime effects are treated as central: bandwidth sharing and compute–communication overlap. For bandwidth sharing, Proteus assumes that concurrent communication operators fairly share available bandwidth. It determines how many communication groups share each physical link in the cluster topology and proceeds hierarchically from NIC bandwidth down through the topology. The model also handles the intersection between feature and gradient communications, adapting only the overlapped portions. For compute–communication overlap, Proteus tracks operator start and end times. If a computation operator overlaps with a gradient communication operator, or conversely, the operation is considered overlapped. Rather than using a learned overlap law, Proteus introduces a fixed overlap factor $\gamma$, whose value is obtained by profiling the increase ratio observed during overlapping backward passes in data-parallel training.

Communication costs are estimated with an $\alpha$-$133.8\%$0 model and then adjusted using topology-aware bandwidth and collective characteristics. The paper additionally notes that the bandwidth of a communication group is computed by summing the relevant communication channels identified via NCCL topology detection. As a point of comparison, it cites an alternative overlap model used in Pollux,

$133.8\%$1

but Proteus itself does not use this formulation. Instead, it adopts the simpler machine- and model-specific $133.8\%$2 factor. This is a deliberate design choice: the simulator favors explicit execution modeling plus lightweight profiled corrections over a more learned overlap estimator.

5. Empirical evaluation

Proteus is implemented as a Python library of about 9K lines of code and follows the PyTorch API. The evaluation covers six representative models: ResNet50, Inception_V3, VGG19, GPT-2, GPT-1.5B, and DLRM. These span vision, NLP, and recommendation workloads. All experiments use synthetic datasets to remove data-loading latency. The evaluation hardware consists of three configurations:

Configuration	Hardware
HC1	1 node, 8 TitanXp GPUs, PCI-e
HC2	4 nodes, 8 V100 GPUs per node, NVLink intra-node, 100 Gbps inter-node
HC3	2 nodes, 8 A100 GPUs per node, NVLink, 200 Gbps inter-node

The software stack is PyTorch 1.8, CUDA 10.1, cuDNN 7.6.5, and NCCL 2.7.8 (Duan et al., 2023).

The main reported result is a $133.8\%$3 average prediction error across the evaluated training throughput results. Among 180 simulation results, out-of-memory status is mispredicted only twice. Proteus is also reported to preserve the ordering of throughput across parallelization strategies, which is important for selecting among candidate strategies. Against a reimplementation of FlexFlow’s simulator adjusted to use collective communication instead of point-to-point, Proteus reports an average error of $133.8\%$4 versus $133.8\%$5, and a maximum error of $133.8\%$6 versus $133.8\%$7, summarized as a reduction in prediction error by up to $133.8\%$8. The same comparison notes that the alternative simulator fails to estimate performance for one-third of the 180 tasks, whereas Proteus handles them.

The paper also provides strategy-level evidence. For GPT-2, Proteus reports a $133.8\%$9 average error while preserving the rank order across multiple data, model, and pipeline combinations. On HC2, it correctly predicts that data-parallel-heavy strategies are best because model parallelism incurs bandwidth sharing on the interconnect and pipeline parallelism introduces bubbles. This is significant because the relative ordering of strategies is often more operationally important than exact throughput values.

An ablation study isolates the contribution of HTAE’s runtime-behavior modeling. Without the runtime behavior detector, average error rises to $\mathcal{P}$0; with the full Proteus system, the reported error is $\mathcal{P}$1 on the tested VGG19 and GPT-2 cases. The paper further states that VGG19 is particularly sensitive to compute–communication overlap, whereas GPT-2 is more sensitive to bandwidth sharing. This result supports the claim that both effects are necessary for accurate simulation rather than marginal refinements to a layerwise cost model.

Simulation overhead is modest. On HC2, total simulation time ranges from $\mathcal{P}$2 s for VGG19 on one GPU to $\mathcal{P}$3 s on 32 GPUs, and from $\mathcal{P}$4 s to $\mathcal{P}$5 s for GPT-2 across the same scale range. The significance of these numbers lies in the fact that strategy evaluation does not consume GPU training resources.

6. Scope, limitations, and place in the distributed-training toolchain

Proteus is best understood as a programmable performance-analysis tool for distributed DNN training. Its Strategy Tree decouples strategy specification from model expression and captures nested combinations of operator- and subgraph-level parallelism. Its execution-graph compilation makes communication insertion, scheduling constraints, and recomputation explicit. Its runtime model targets the two phenomena the paper identifies as most important for modern distributed training performance: compute–communication overlap and bandwidth sharing (Duan et al., 2023).

The paper also defines clear boundaries for the system. Proteus currently focuses on commonly used communication primitives and would require extension to support new ones. It models compute–communication overlap and bandwidth sharing, but not arbitrary overlap among multiple computation operators on a single GPU. Its accuracy depends on operator-cost estimators or profiling inputs, as well as on accurate topology descriptions for communication modeling. The evaluation uses synthetic data rather than end-to-end input pipelines. Finally, Proteus is aimed at modeling specified strategies rather than searching over the entire parallelization space, so it is complementary to automated parallelization frameworks rather than a replacement for them.

A common misconception would be to interpret Proteus as a generic black-box performance oracle for all aspects of training. The paper does not support that reading. Its contribution is narrower and more technical: it provides a simulation-based framework for distributed throughput prediction under complex parallelization, with explicit treatment of redistribution, micro-batch scheduling, memory release, topology-aware communication, and overlap effects. Within that scope, the reported results suggest that accurate standalone simulation can be practical enough to inform strategy comparison, bottleneck analysis, and resource planning before deployment.