Graph Partitioning

The primary objective is to distribute computational work evenly across available processors or devices. An imbalanced partition leaves some units idle while others are overloaded, creating a bottleneck.

• Goal: Minimize the makespan (total execution time) of the slowest partition.
• Challenge: Operations have varying FLOP costs and memory footprints.
• Example: Partitioning a Vision Transformer across 4 GPUs requires balancing the compute-heavy self-attention layers, not just splitting layers evenly.

Partitions must exchange data across boundaries. This communication is expensive in terms of latency and bandwidth.

• Goal: Minimize the edge cut—the number or volume of tensors passed between partitions.
• Strategy: Use algorithms like Kernighan–Lin or METIS that cut few edges.
• Impact: High communication can negate parallel speedup, adhering to Amdahl's Law. Fusing adjacent operators before partitioning is a common tactic to reduce edges.

Partitions must respect the capabilities and constraints of the target hardware. This is not a purely mathematical graph problem.

• Constraints: Memory capacity per device, supported data types (e.g., INT8), and specialized accelerator support (e.g., NPU-compatible ops).
• Strategy: A cost model estimates latency for a subgraph on a specific device. The partitioner assigns subgraphs to the hardware where they execute fastest.
• Example: Offloading convolutional layers to a DSP while running element-wise ops on a CPU.

For models too large for data parallelism, partitioning enables pipeline parallelism, where different stages of the model run concurrently on different devices.

• Mechanism: The graph is split at layer boundaries. While Device 2 processes micro-batch N, Device 1 processes micro-batch N+1.
• Challenge: Avoiding pipeline bubbles (idle time) requires careful scheduling (e.g., 1F1B schedule).
• Use Case: Essential for large language model inference, splitting transformer blocks across multiple devices.

When a single layer's parameters exceed a device's memory, model parallelism splits the layer itself across devices. Graph partitioning identifies these intra-layer split points.

• Tensor Parallelism: Splits weight matrices (e.g., in a feed-forward layer) across devices, requiring all-reduce communication after the operation.
• Example: A 120B parameter dense layer is split across 8 GPUs, each holding 15B parameters. The partitioner must insert the correct communication collectives.

Grouping operations that reuse the same data into the same partition improves cache hit rates and reduces expensive DRAM accesses.

• Principle: Follows the same logic as loop tiling in traditional compilers.
• Benefit: Especially critical for memory-bound operations and on devices with complex memory hierarchies (e.g., CPU with L1/L2/L3 cache).
• Technique: Algorithms that maximize modularity in the graph naturally group tightly-connected operations.

Frameworks analyze the computational graph and available hardware to assign subgraphs to the most suitable processor. This involves a cost model that estimates the execution latency of each operator on different backends (e.g., CPU, GPU, NPU).

• Key Strategy: Operators with high arithmetic intensity (e.g., large matrix multiplications) are assigned to parallel accelerators like GPUs, while control-flow-heavy or memory-bound ops may run on CPUs.
• Example: In TensorFlow, a Hardware Delegate (e.g., TfLiteGpuDelegate) is responsible for claiming and executing a supported subgraph on the target accelerator.
• Outcome: Maximizes aggregate hardware utilization and minimizes data transfer between devices.

The graph is partitioned depth-wise, splitting consecutive layers or stages across different devices. This enables pipeline parallelism, where different devices process different micro-batches simultaneously.

• Mechanism: The framework inserts send and receive operators at partition boundaries to handle inter-device communication.
• Use Case: Essential for large models that do not fit into the memory of a single device. Each device holds a portion of the model's weights.
• Framework Example: PyTorch's Pipe API in torch.distributed.pipeline.sync implements this by splitting a nn.Sequential module across GPUs.

Individual weight matrices within a layer (e.g., a linear transformation) are split across multiple devices. Each device computes a portion of the output, requiring a communication step (like an all-reduce) to combine results.

• Target: Extremely large layers, common in transformer models with multi-billion parameter feed-forward or attention blocks.
• Implementation: Frameworks like Megatron-LM split the weight matrix of a linear layer along its column or row dimension. The computation Y = XA is distributed by splitting matrix A.
• Benefit: Allows training and inference of layers whose parameters exceed the memory of a single accelerator.

The compiler performs a graph traversal (often a topological sort) to analyze data dependencies and find optimal cut points. The goal is to minimize the cost of cross-partition communication while balancing computational load.

• Algorithm: Uses heuristics or a cost model to evaluate potential partition boundaries. The communication cost is modeled as the size of tensors that must be transferred.
• Tooling: XLA's partitioner and TVM's relay.transform.PartitionGraph pass perform this analysis, considering user-provided compiler hints (@tvm.register_func) for external code generation.
• Constraint: Must preserve the original graph's semantic correctness; all producer operations must execute before their consumers.

After deciding on partition boundaries, the compiler automatically inserts the necessary communication operators into the graph.

• Primitives Include: Send/Recv, AllReduce, AllGather, and Broadcast. The choice depends on the parallelism strategy.
• Location: Inserted at the edges where tensors flow from one partition to another.
• Optimization: Frameworks may fuse these communication ops with neighboring computation or schedule them asynchronously to overlap communication with computation, hiding latency.

The partitioned subgraphs are managed by a runtime scheduler that handles execution order, synchronization, and memory transfer between devices.

• Role: Ensures data dependencies across partitions are satisfied. Manages device-specific operator dispatch.
• Example: TensorFlow's Placer assigns nodes to devices, and its Executor schedules them. For multi-device graphs, the CollectiveOps are scheduled to coordinate communication.
• Advanced: Some runtimes perform dynamic re-partitioning based on real-time load or input data characteristics, though static partitioning at compile-time is more common for predictability.

A cost model in a compiler is an analytical or heuristic function that estimates the computational expense of executing an operation, subgraph, or schedule. It is fundamental for making intelligent graph partitioning decisions.

• Purpose: Guides the partitioner by predicting metrics like latency, memory bandwidth usage, and power consumption for a candidate partition on a specific processor.
• Implementation: Can be static (based on operator type and tensor shapes) or dynamic (profiled from previous runs).
• Use Case: A partitioner uses a cost model to evaluate whether fusing two operators or splitting them across a CPU and NPU yields lower predicted latency.

A topological sort produces a linear ordering of the nodes in a Directed Acyclic Graph (DAG) where for every directed edge from node A to node B, node A appears before node B. This is the foundational algorithm for scheduling operations within a partitioned subgraph for execution.

• Prerequisite: Graph partitioning outputs subgraphs that are themselves DAGs.
• Runtime Scheduling: The topological order within each partition defines the sequence of kernel launches, ensuring data dependencies are respected.
• Parallelism: While the sort provides a valid serial order, it also reveals independent operations that can be executed in parallel on available threads or streams.

Static memory planning is a compile-time optimization that pre-allocates and reuses memory buffers for tensors by analyzing their lifetimes within a computational graph. It is crucial for efficient execution of partitioned subgraphs, especially on memory-constrained devices.

• Process: The compiler analyzes the liveness of each tensor (when it is first produced and last consumed) to overlap memory for tensors with non-overlapping lifetimes.
• Benefit: Minimizes peak memory footprint and eliminates the overhead of dynamic memory allocation during inference.
• Partitioning Context: When a graph is partitioned, memory planning can be performed per-subgraph or globally across partitions to manage data transfer buffers between devices.

Operator dispatch is the runtime mechanism that selects the optimal kernel implementation for an operation based on runtime parameters. In a partitioned graph, efficient dispatch is key to performance within and across hardware boundaries.

• Dispatch Factors: Input data types (float32, int8), tensor shapes, and available hardware (CPU core, GPU, NPU).
• Multi-Backend Execution: For a partition targeting a heterogeneous system (e.g., CPU+GPU), the dispatch layer routes each operator to the correct backend's kernel library.
• Performance Impact: A high-overhead dispatch mechanism can negate the benefits of graph partitioning; thus, it's often optimized via pre-compiled dispatch tables.

The Roofline model is an intuitive visual performance model used to analyze the maximum achievable performance of a computational kernel or subgraph. It informs graph partitioning by characterizing the hardware's limits and the partition's inherent characteristics.

• Axes: Plots performance (GFLOPs/sec) against operational intensity (operations per byte of DRAM access).
• Two Limits: Performance is bounded either by peak compute throughput (the 'roof') or by memory bandwidth (the 'slope').
• Partitioning Insight: A partitioner can use this model to identify if a subgraph is compute-bound (benefits from an accelerator) or memory-bound (may benefit from CPU cache optimizations), guiding device placement decisions.