Source-linked AI summary
Supporting Very Large Models using Automatic Dataflow Graph Partitioning
Minjie Wang, Chien-chin Huang, Jinyang Li
TL;DR
Very large DNNs are constrained by limited GPU memory, motivating transparent partitioning across multiple devices. Tofu describes fine-grained operator semantics in TDL and recursively searches dataflow-graph partitions to minimize communication. On eight GPUs, it trains models that do not fit on one GPU and achieves 25% - 400% higher throughput than alternative approaches.
Problem
Limited GPU memory constrains the size of DNN models that can be explored, motivating partitioning across multiple devices.
Method
Tofu analyzes TDL descriptions of fine-grained operators and uses recursive graph search to find partition plans that minimize total communication.
Results
25% - 400% higher training throughput was achieved than alternative approaches on large models using an eight-GPU machine.
Takeaways & Limitations
Tofu enables training very large DNN models across multiple GPUs while working with a general-purpose dataflow platform.
Takeaways & Limitations
Tofu is restricted to partition-n-reduce, does not exploit interconnect topology, and cannot automatically partition some operators or sparse tensor operations.
Abstract
from arXiv · showhide
This paper presents Tofu, a system that partitions very large DNN models across multiple GPU devices to reduce per-GPU memory footprint. Tofu is designed to partition a dataflow graph of fine-grained tensor operators in order to work transparently with a general-purpose deep learning platform like MXNet. In order to automatically partition each operator, we propose to describe the semantics of an operator in a simple language which represents tensors as lambda functions mapping from tensor coordinates to values. To optimally partition different operators in a dataflow graph, Tofu uses a recursive search algorithm that minimizes the total communication cost. Our experiments on an 8-GPU machine show that Tofu enables the training of very large CNN and RNN models. It also achieves 25% - 400% speedup over alternative approaches to train very large models.
1 Introduction
Tofu addresses GPU-memory limits by automatically partitioning fine-grained tensor operators across devices on general-purpose dataflow platforms. It uses TDL to analyze operator semantics and recursive search to optimize graph-wide partitioning, achieving higher training throughput on very large models while retaining important limitations.
- Motivation: GPU memory constrains the size of DNN models that can be explored despite continuing growth in model parameters.The number of parameters in state-of-the-art neural networks has roughly doubled every 2.4 years since the 1980s.
- Motivation: Model parallelism reduces per-GPU memory footprint and can provide parallel speedup, with tensor partitioning preferred for very large models.Tensor partitioning balances per-GPU memory usage and is necessary for speeding up popular CNN models.
- Challenge: Existing tensor-partitioning methods operate mainly at coarse layer granularity, requiring specialized implementations or restricting models to common layer compositions.This limits their applicability to general-purpose dataflow systems containing many fine-grained tensor operators.
- Approach: Tofu automatically partitions fine-grained operator inputs and outputs in MXNet, with a solution that could potentially apply to other dataflow systems.This operator-partitioning approach is more fine-grained than layer partitioning.
- Approach: TDL describes operator semantics as output values derived from input tensors, enabling symbolic analysis of required input regions and graph-wide communication planning.Tofu also shrinks the search space through graph coarsening and recursive two-worker partitioning.
- Results: 25% - 400% higher training throughput was achieved than alternative approaches on large Wide ResNet and recurrent models using eight GPUs.Most evaluated models did not fit in a single GPU’s memory.
- Limitations: Tofu cannot automatically partition some operators, does not exploit communication topology, and is intended for very large rather than moderately sized models.The paper identifies these limitations as requiring further research.
2 Background
The background motivates transparent tensor partitioning across multiple GPUs as a way to train models that exceed individual device memory. The approach must also account for communication costs because GPU interconnect bandwidth can limit performance.
- Memory constraints: GPU device memory is smaller than CPU memory, ranging from 12GB on NVIDIA K80s to 16GB on NVIDIA Tesla V100s.Google TPU cores have a similar limitation, with 8GB attached to each core.
- Tensor partitioning: Partitioning tensors across k devices roughly reduces each device’s computation memory requirement to 1/k of single-device memory.Partitioning also provides performance speedup through parallel execution.
- System goal: Tofu aims to automatically partition tensors and parallelize operators in a dataflow graph transparently to users.The target platforms include TensorFlow and MXNet.
- Communication constraint: The communication required by partitioned execution can exceed deployed GPU-cluster network bandwidth, making communication a central system constraint.The required aggregate bandwidth is determined by transferred bytes divided by computation time.
3 Challenges and our approach
Tofu tackles automatic partitioning at operator and graph levels. It analyzes TDL-described operators under partition-n-reduce and uses search-space reduction with recursive optimization for graph-wide partition plans.
- 3.1 How to partition a single operator?: Tofu uses partition-n-reduce, where workers run the original operator on smaller inputs and combine outputs by concatenation or element-wise reduction.This pattern reuses existing optimized single-GPU implementations.
- 3.1 How to partition a single operator?: Partition-n-reduce is broadly useful but cannot express every parallel algorithm and may communicate more than specialized methods.Cholesky is cited as an inapplicable example, while Cannon’s algorithm can achieve more efficient matrix-multiplication communication.
- 3.1 How to partition a single operator?: Conv1d admits multiple partition strategies, including concatenating outputs across the batch dimension or reducing worker outputs across a reduction dimension.Different strategies require different input tensor regions and remote data fetches.
- 3.1 How to partition a single operator?: Prior methods manually discover strategies for only a few common layers, a costly approach for systems with 341 TensorFlow or 139 MXNet operators.Some methods also omit output-reduction strategies that can improve performance.
- 3.1 How to partition a single operator?: TDL lets developers provide separate, high-level descriptions of operator computation so Tofu can analyze access patterns and discover viable partition strategies.The description is separate from the operator implementation and abstracts away algorithmic architecture.
- 3.2 How to optimize partitioning for a graph?: Optimal partitioning of a general dataflow graph is NP-hard, and fine-grained tensor choices cause the number of strategies and search time to grow rapidly with GPU count.Each tensor may be partitioned along combinations of multiple dimensions.
- 3.2 How to optimize partitioning for a graph?: Tofu makes graph search practical by coarsening related operators and recursively applying a dynamic-programming search to limit the number of workers considered at each step.Coarsening groups forward and backward operations and coalesces element-wise or unrolled operators.
4 Partitioning a single operator
Tofu uses TDL to describe tensor operators as lambda functions, then applies symbolic interval analysis to discover viable partition strategies and required input regions. These strategies guide later graph partitioning and remote data fetching.
- TDL: TDL represents tensors as lambda functions mapping coordinates to values through side-effect-free expressions and reductions.Supported constructs include index variables, tensor elements, arithmetic operations, and reductions such as Sum, Max, Min, and Prod.
- Limitations: TDL uses an opaque-function primitive for computations it cannot express, such as Cholesky decomposition.A batched opaque operator may still be partitioned along its batch dimension.
- TDL: TDL can describe 134 of 139 MXNet v0.11 operators, including element-wise operators, opaque functions, and output reductions.The descriptions were written to bootstrap Tofu for an existing MXNet dataflow system.
- Partition strategies: Tofu uses the inferred input regions to specify each worker’s share of an operator and to generate the partitioned graph’s remote data fetches.The same region information supports both graph-level optimization and partitioned graph construction.
- Analyzing TDL: Symbolic interval analysis executes an operator’s lambda expression to infer input tensor access ranges for symbolic output-index bounds.This avoids repeatedly analyzing thousands of operators with concrete tensor shapes.
- Partition strategies: The analysis discovers both output-dimension partitioning and partition-n-reduce strategies, including alternatives for convolution output and reduction dimensions.For conv1d, output partitioning can read partitioned data with full filters, while reduction partitioning produces partially reduced tensors.
5 Partitioning the dataflow graph
Tofu coarsens fine-grained dataflow graphs so dynamic programming can optimize them, then recursively applies two-group partitioning to support many workers and multiple tensor dimensions. The objective is minimizing total communication cost, which also reduces communication-buffer memory.
- Optimization goal: Tofu minimizes total communication cost because kernel time is relatively insensitive to partition dimensions and communication buffering contributes to per-worker memory.The tensor-share portion of memory is fixed at 1/k of single-device memory for k GPUs.
- Graph coarsening: To make non-linear operator graphs suitable for dynamic programming, Tofu coarsens them by grouping or coalescing operators and tensors.Coarsened graphs for MLPs and CNNs become linear.
- Graph coarsening: Forward operators are grouped with generated backward operators, while forward tensors are grouped with their gradients and any required gradient summation.This preserves the relationship between training computations during graph optimization.
- Graph coarsening: Tofu also coalesces consecutive element-wise operators and unrolled RNN timesteps when they should share a partition strategy.RNN timesteps share computation logic and weights, allowing a multi-layer RNN graph to become a chain of coalesced and grouped operators.
- Recursive partitioning: Recursive partitioning applies the two-worker dynamic-programming plan repeatedly, enabling multi-dimensional tensor partitions across k = 2^m GPUs.Each recursive step partitions worker groups and carries fetched data into subsequent subproblems as extra inputs.
- Recursive partitioning: The recursive partitioning theorem states that total communication cost across worker groups does not decrease across recursion steps.For step costs δ_i, the stated result is δ_i ≤ δ_i+1.
- Recursive partitioning: Recursive search reduces the strategy count for 8-worker convolution from 2^06 to 3 ∗ 4096 and searches plans compatible with hierarchical interconnects.The recursion assigns lower-communication worker groups near slower, higher-level links earlier in the hierarchy.
6 Optimizations in generating the partitioned graph
After selecting partition strategies, Tofu generates a partitioned dataflow graph and adds optimizations that preserve memory reuse, streamline remote fetches, distribute reductions, and delay fetch execution.
- Memory reuse: Naive partitioned graph generation can increase per-worker memory because changed dependencies prevent existing memory planners from immediately reusing buffers.Tofu addresses this by preserving original operator dependencies on each worker with extra control dependencies.
- Remote data fetch: Tofu generates worker-local copies of operators and uses data-movement operators to fetch, copy, and assemble remote input regions.The implementation uses MXNet copy, split, and concatenate operators for remote data movement.
- Execution optimizations: Tofu spreads output-reduction work across all GPUs and delays remote-fetch execution to avoid aggregation bottlenecks and unnecessarily prolonged memory occupancy.The reduction optimization uses all-reduce, while fetch scheduling follows a technique adopted from TensorFlow.
7 Evaluation
Tofu trains very large WResNet and RNN models across eight GPUs, generally approaching ideal throughput and outperforming alternative memory-management and placement strategies. Its recursive partitioning also finds complex, multidimensional plans that reduce communication and improve throughput.
- Training large and deep models: Tofu trains very large WResNet and RNN models across 8 GPUs at 60%-98% of ideal throughput.WResNet reaches 60%-95%, while RNN configurations reach 70%-98% of ideal throughput.
- Training large and deep models: 25%-400% higher training throughput is achieved by Tofu than by other approaches for large DNN models.The evaluation compares Tofu with reduced-batch, CPU-swapping, and operator-placement alternatives.
- Comparing partition algorithms: Tofu’s partition algorithm outperforms competing algorithms and heuristics, while multidimensional partitioning and output reduction avoid failures on demanding models.AllRow-Greedy runs out of memory on WResNet-152-10, and ICML18 does so without output reduction.
- Training large and deep models: Swapping is 20%-63% slower than Tofu across all WResNet models because CPU-GPU communication becomes the bottleneck.All eight GPUs share the same bandwidth when communicating with CPU memory, even with prefetching.
- Training large and deep models: Operator placement achieves 38%-61% of Tofu’s throughput and cannot train RNN-10-8K because layer-wise pipelining leaves GPUs underutilized.Tofu instead parallelizes each operator and keeps all GPUs busy at all times.
- Partition results: Tofu partitions tensors across batch and channel dimensions, combining different strategies across convolution layers and switching fetched tensors by layer.It fetches weights for lower layers but relatively smaller activations for higher layers, reflecting their changing tensor sizes.
8 Related Work
Prior systems support parallel tensor computation, but often require specialized operators or coarse-grained layer compositions. Tofu instead automates fine-grained operator partitioning and whole-graph partition optimization using TDL.
- Unlike data-parallel training, model parallelism distributes parameters across GPUs and is suitable for very large models.
- Model compression reduces model size through pruning, quantization, or reduced precision, but can affect accuracy; Tofu preserves model behavior while enabling larger models.
- Existing parallel tensor libraries and programming systems provide efficient computation or high-level primitives, but adding new operators can require substantial manual effort.
- Earlier automatic parallelization systems specialize in tensor contractions, map/reduce primitives, or array operators rather than general dataflow graphs.
- Tofu automatically discovers partition-n-reduce patterns from TDL descriptions and optimizes partitioning across the entire dataflow graph.
9 Discussion, limitations, and future work
Tofu’s partition-n-reduce design and TDL language impose important scope and optimization limits. The system cannot cover some operators, fully exploit hardware topology, or flexibly adapt partitioning to model and device heterogeneity.
- Tofu’s partition-n-reduce scheme cannot parallelize every computation, including Cholesky, and does not necessarily minimize communication or exploit interconnect topology.
- TDL lacks control flow and data-dependent indexing, while Tofu does not support sparse tensor operations because of load imbalance.
- Tofu does not verify that operator implementations match their TDL descriptions, leaving verification as an open research problem.
- Tofu always partitions every operator and tensor across all workers, which can produce undersized GPU kernels for moderately sized models.
- Tofu lacks support for non-uniform partitioning across GPUs with different computing or memory capacities and does not explicitly optimize for interconnect topology.
- Its dynamic-programming search cannot optimize device placement choices for unpartitioned or non-uniformly partitioned operators; stochastic search is future work.
10 Conclusion
Tofu enables training very large DNN models by partitioning dataflow graphs across GPUs. It automates operator strategy discovery from TDL and searches for a communication-minimizing whole-graph plan.
- Tofu enables training very large DNN models by partitioning a tensor dataflow graph across multiple GPU devices.
- Tofu infers valid operator partition strategies from semantics written in its lightweight TDL description language.
- Tofu combines dynamic programming with DNN-specific heuristics to find a partition plan minimizing communication across the entire dataflow graph.
A.1 Recursive partition plan
Tofu represents a dataflow partition plan through recursive one-dimensional splits across worker groups. Applying each basic plan to all resulting subgraphs yields a sequence of partition decisions over the GPUs.
- A partition plan specifies how each tensor is partitioned and how each operator is parallelized, with splits potentially applied along multiple tensor dimensions.
- For 2^m GPUs, any plan can be represented as m recursive basic plans, each splitting tensors along one dimension between two worker groups.
- After i recursive steps, the graph becomes 2^i identical subgraphs whose tensors are 1/2^i the original size, and the next plan applies to every subgraph.
- Tofu searches for a sequence of recursive partition plans that is no worse than the optimal sequence under its analysis.
A.2 Region Analysis
The analysis characterizes operator access patterns and communication costs under restricted affine-indexing assumptions. It then establishes commutativity of basic partition plans and proves the recursive algorithm optimal.
- Region Analysis: Symbolic intervals represent output-index bounds and input access ranges, producing an affine transformation for operator access analysis.X_i bound output indices, while Y_i describe input-dimension access ranges.
- Assumptions: The analysis restricts operators so each output index accesses at most one dimension per input tensor and input dimensions scale linearly with one output index.These assumptions exclude examples such as A[i, i] and partition-n-reduce strategies such as convolution halo exchange.
- Region Analysis: Under these assumptions, an input tensor shape is represented as β1Xπ1 × . . . × βdXπd, with constant factors and a permutation of output dimensions.This corollary provides the shape structure used in communication analysis.
- Communication Cost: Communication cost is a weighted sum of tensor sizes because partitioning may require fetching unavailable input regions or sending output regions to other devices.The total cost aggregates within-group communication over the tensors in the dataflow graph.
- Optimality Proof: The recursive algorithm is optimal because any sequence with a larger total cost would permit a lower-cost next partition, contradicting dynamic programming per-step optimality.The contradiction compares an allegedly worse sequence with an optimal sequence after reordering their per-step costs.