2000 character limit reached
Piper: A Programmable Distributed Training System
Published 9 Jun 2026 in cs.DC and cs.AI | (2606.11169v1)
Abstract: Large-scale model training increasingly relies on composing multiple parallelism strategies, such as data, pipeline, and expert parallelism, together with memory-saving optimizations like ZeRO. Deployed systems for foundation model pretraining often rely on human experts to manually design a high-level parallelism strategy then implement the corresponding low-level execution strategy, making it difficult to adapt the system to new strategies. Meanwhile, many general-purpose frameworks are more flexible but their implementations are still tied to a fixed set of common parallelism strategies, making it challenging to integrate state-of-the-art strategies. We present Piper, a user-controllable distributed training system that decouples the strategy from the runtime implementation. Piper allows users to declare a comprehensive distributed training strategy with a small set of model annotations and scheduling directives. Each directive applies a transformation on Piper's intermediate representation (IR), a unified global training DAG that represents all computation and communication. Using this IR, Piper compiles per-device execution plans and executes them with a distributed runtime agnostic to the strategy. We show that the combined system maintains performance parity on commonly available strategies such as ZeRO, while also enabling additional performance and memory efficiency gains through joint scheduling of compute and communication in composed parallelism strategies such as DeepSeek-V3's DualPipe.
Paper Prompts
Sign up for free to create and run prompts on this paper using GPT-5.
Top Community Prompts
Explain it Like I'm 14
Overview
This paper introduces Piper, a system that helps train very large AI models across many computers (GPUs) more easily and efficiently. Today, experts often have to hand-craft complicated training “plans” to make big models run fast on big clusters. Piper lets users describe what plan they want in a simple, programmable way, and then it automatically figures out how to run it efficiently on the hardware.
What questions were the authors asking?
- Can we make it easy for people to describe any training plan they want (for example, how to split work across devices and when to communicate), without rewriting lots of low-level code?
- Can we turn that plan into fast, reliable execution on many GPUs, even for new, advanced strategies that current tools struggle to support?
- Can one system both match standard methods on common strategies and also unlock more speed and memory savings when strategies are combined?
How did they do it?
The big idea: separate the plan from the engine
Think of training as a road trip:
- The “plan” is the route you choose (which roads, when to stop, where to meet).
- The “engine” is the car that actually drives.
Most systems mix the two, so changing the route means rebuilding parts of the car. Piper separates them. You describe the route (the training strategy), and Piper’s engine makes the car follow it efficiently.
A common “map” for everything (a global DAG IR)
Piper builds one big flowchart of the whole training step called an IR (Intermediate Representation). It’s a directed acyclic graph (DAG), which you can think of as a step-by-step map:
- Boxes (nodes) are tasks, like “do this layer’s math” or “send data to another GPU.”
- Arrows show what must happen before what.
- Every computation and every communication is on this map, so Piper can schedule them together smartly.
Two kinds of nodes:
- Chunk: a block of compute (like doing a layer forward or backward).
- Comm: a communication step (like “all-reduce” to average gradients or “all-to-all” to shuffle expert data).
Simple tools (“directives”) to shape the plan
You lightly label your model (with “annotations”) to mark meaningful parts, like pipeline stages or expert layers. Then you apply a few simple commands (directives) to transform the map. These are like giving instructions to a team:
- Place: put a part of the model on certain devices (who does what).
- Replicate: make copies of a part on many devices to do the same work on different data, then sync results (this is data parallelism; sync is usually an “all-reduce”).
- Shard: split weights or work across devices (like expert parallelism); adds “all-to-all” steps to reassemble the right pieces.
- Split: cut a batch into microbatches so different stages can overlap like an assembly line (pipeline parallelism).
- Order: set the execution order, including overlaps (e.g., run forward of microbatch A and backward of microbatch B at the same time to hide communication delays).
You can also choose which “stream” to use. A stream is like a lane on a highway—separate lanes let tasks run side-by-side without blocking each other.
Compiling and running
- Compiler phase:
- Trace the model to get its operator graph.
- Cut it into “Chunks” where you placed annotations.
- Apply your directives to insert the right Comm nodes (like all-reduce or all-to-all), split into microbatches, assign devices and streams, and set ordering.
- Clean up redundant communications (e.g., reuse a previous gather when possible).
- Runtime phase:
- A central scheduler creates a per-device “to-do list,” ordered by streams and dependencies.
- Each worker (GPU) executes its list, managing:
- Streams (lanes) to overlap work.
- Communicators (channels) for collective ops without deadlocks (they even separate send and receive into different channels to avoid traffic jams).
- Memory (like ZeRO, which shares and rematerializes weights/gradients to save space).
The runtime is “strategy-agnostic,” meaning it doesn’t care which plan you chose—it just follows the map.
A simple example: DualPipe, explained
DualPipe is an advanced schedule from DeepSeek-V3. It overlaps forward on one microbatch with backward on another to hide slow communication (especially expert “all-to-all”). In normal tools, fitting this neatly can be hard. In Piper, you:
- Place pipeline stages on devices.
- Replicate non-expert layers for data parallelism.
- Shard expert layers for expert parallelism.
- Split into microbatches.
- Order them to overlap forward and backward.
- Assign streams so compute and communication don’t block each other.
Piper compiles this into a global map and runs it efficiently.
What did they find?
- Piper matches the speed of common frameworks (like Megatron and TorchTitan) on standard strategies (e.g., ZeRO for data parallelism).
- On combined/advanced strategies (like DualPipe mixing pipeline and expert parallelism with smart overlap), Piper delivers 6–30% higher throughput than those baselines.
- Piper can support combinations that other general frameworks struggle with (for example, pipeline parallelism plus different ZeRO levels), enabling 3–8× larger batch sizes. Larger batches can mean better hardware use and faster training.
Why this matters:
- Combining strategies can cause hidden slowdowns (network contention, memory pressure). Piper’s unified map and scheduler coordinate compute and communication together, reducing these problems.
Why does this matter?
- Faster training for big models: More overlap and smarter scheduling mean less waiting and more doing.
- Easier experimentation: Researchers can try new strategies by describing them with a few annotations and directives, instead of writing complex, low-level code.
- Future-proof: As new ideas appear (new kinds of parallelism or schedules), users can express them in Piper without changing the runtime.
- Better hardware use: Jointly scheduling compute, communication, and memory helps avoid bottlenecks and fits bigger batches/models on the same GPUs.
Key terms in simple words
- Data parallelism (DP): Many copies of the model process different data, then average their updates (like students solving different worksheets and then comparing answers).
- ZeRO: A memory-saving version of DP that splits model states (weights/gradients/optimizers) across devices and only gathers them when needed (like sharing heavy textbooks and borrowing them only when you need a page).
- Pipeline parallelism (PP): Split the model into stages and feed microbatches like an assembly line, so stages work in parallel.
- Expert parallelism (EP): In Mixture-of-Experts models, different GPUs hold different “experts.” Inputs are routed to the right experts, then results are combined (this requires “all-to-all” communication).
- All-reduce: A group operation to average or sum values across GPUs.
- All-to-all: Every GPU sends some data to every other GPU (used for routing to experts).
- Stream: A lane of work on a GPU; multiple lanes let tasks run at the same time.
- DAG (Directed Acyclic Graph): A flowchart with arrows showing what must happen before what; no cycles means you don’t loop back.
Knowledge Gaps
Knowledge gaps, limitations, and open questions
Below is a consolidated list of concrete gaps and open problems that the paper leaves unresolved, framed to guide future research and engineering work:
- Limited sharding generality: Shard currently assumes weight sharding along dimension 0 and specific EP/DP compositions; there is no support for arbitrary tensor sharding specifications, non-uniform shard sizes, or general GSPMD-style inference and propagation of placements across the DAG.
- Incomplete tensor placement inference: The compiler requires explicit device and stream assignments and rejects missing placements; automatic propagation/inference of placements and streams (beyond simple cases) is future work.
- Tied/shared parameters across placements: A single model-state bucket cannot span Chunks with different placements (e.g., tied embeddings across PP stages); mechanisms for correct cross-placement gradient synchronization and optimizer updates are not provided.
- Dynamic models not supported: Piper requires fully traceable models; support for dynamic control flow, data-dependent shapes, Python-side branching, custom CUDA ops, and mixed eager/compiled execution within a single distributed DAG is absent.
- PP schedule ergonomics and automation: Users must specify Order constraints (and nested overlaps) manually; there is no DSL, validator, or auto-scheduler/cost model to synthesize and verify PP schedules for heterogeneous models.
- Safety and correctness guarantees are informal: Deadlock freedom relies on schedule properties (e.g., consistent microbatch ordering) and stream/communicator choices; there is no formal proof or runtime checker/repair for incorrect Order rules, circular temporal dependencies, or mismatched collective orders.
- Heuristic scheduling without performance model: The centralized and worker-level schedulers use simple graph-degree and hand-written priority rules with no compute/comm cost model, no profile-guided optimization, and no online adaptation to workload or network variability.
- Centralized scheduler scalability and resilience: A single centralized scheduler can be a throughput and fault-tolerance bottleneck at hundreds/thousands of GPUs; there is no analysis of its latency, memory footprint, or a distributed/hierarchical alternative.
- Communicator scaling limits: Allocating separate communicators per stream/direction can exhaust NCCL limits, add startup overhead, and increase memory usage; policies for communicator pooling, reuse, or consolidation are not explored.
- Topology-agnostic placement and scheduling: The system does not account for NVLink/PCIe/InfiniBand topology, cross-node bandwidth asymmetry, or NUMA effects; topology-aware device mapping and stream/comm scheduling are open.
- Interference-aware comm bucketing: Bucket sizes for allreduce/allgather are user-chosen; there is no automatic bucketing or adaptive partitioning to mitigate A2A–AR interference while preserving bandwidth efficiency.
- Multi-tenant and shared-resource interference: The runtime assumes exclusive access to GPU/PCIe/NIC; strategies for fairness and performance isolation under multi-tenancy or background jobs are not addressed.
- Fault tolerance and elasticity: Recovery from worker failures, communicator reinitialization, rescheduling of in-flight DAG segments, elastic scale-up/down, and checkpoint integration are unspecified.
- Determinism and reproducibility: While collective launch order is kept deterministic, cross-stream asynchronous dispatch and local ready-queue choices can make full-run determinism fragile; there is no mechanism to guarantee or test determinism end-to-end.
- Memory management under pressure: Interaction with the PyTorch allocator and fragmentation is only discussed qualitatively; there is no memory-pressure-aware scheduling, defragmentation strategy, or analysis under extreme HBM pressure.
- Activation checkpointing and offloading: Integration with activation rematerialization, CPU/NVMe offloading, or CUDA Unified Memory to expand effective batch/model size is not presented.
- Optimizer and ZeRO coverage: The system demonstrates ZeRO-2 gradient sharding but not full ZeRO-3 optimizer/parameter sharding, optimizer offloading, fused/8-bit optimizers, or correctness of overlapping optimizer updates with comm/compute.
- Mixed precision and low-precision comms: There is no discussion of AMP/grad-scaling per Chunk, FP8/bfloat16 kernels, or reduced-precision collective communication and their scheduling/precision trade-offs.
- MoE dynamics and load imbalance: Static DAG compilation does not address dynamic expert routing skew, variable A2A payloads, capacity overflow, or adaptive expert rebalancing; how to safely adapt schedules at runtime is open.
- Context/sequence parallelism: While mentioned in background, directive support and evaluation for CP/sequence-parallel sharding of activations (and its interactions with PP/EP/ZeRO) are not demonstrated.
- Non-transformer and multimodal workloads: Validation is focused on transformer/MoE; applicability to CNNs, graph neural nets, diffusion models, multimodal encoders/decoders, and RL training loops remains untested.
- Autograd and aliasing corner cases: Correctness under in-place ops, view/alias semantics, custom autograd Functions, and complex gradient flows is not analyzed.
- Kernel fusion and backend interactions: How Piper’s chunk boundaries interplay with TorchInductor/nvFuser fusion, scheduling of fused kernels, and cross-pass optimization conflicts is unclear.
- Ray-based runtime overhead: The impact of Ray’s control-plane and object-store overheads, latency jitter across nodes, and scalability at large GPU counts is not quantified.
- Intra-GPU resource partitioning: The system delegates SM/HBM/NIC contention to the GPU; explicit policies (e.g., SM capping, stream priorities, CUDA graphs) to regulate intra-GPU sharing are not implemented.
- Numerical equivalence and convergence: Reordering via overlaps and stream interleaving could affect numerical trajectories; there is no study of training stability, convergence speed, or accuracy under different schedules.
- Developer tooling and diagnostics: There is no mention of tools for DAG visualization, schedule validation, deadlock detection, comm-order mismatch debugging, or automated reports on bottlenecks and memory hotspots.
- API constraints on filters: Filters must match contiguous sub-DAGs; strategies to relax this constraint, automatically partition non-contiguous matches, or provide actionable compiler errors are not described.
- Generalization beyond NCCL: Portability to alternative backends (RCCL, MPI, UCX, SHARP, TPU/collectives), networks (RoCE vs IB), and heterogeneous accelerators is not addressed.
- Inference and serving: The framework targets training; extension to low-latency/throughput inference, speculative decoding, or streaming generation with similar programmability is unexplored.
Practical Applications
Immediate Applications
Below are practical, deployable-now use cases derived from Piper’s IR, scheduling directives, and runtime, grouped by sector. Each item notes potential tools/products/workflows and key assumptions/dependencies that affect feasibility.
- Industry (Software/AI platforms): Rapid implementation of complex distributed training strategies (e.g., DualPipe-style PP+EP+DP with ZeRO levels)
- What: Use Piper’s annotations and directives (Place, Replicate, Shard, Split, Order) to reproduce and extend strategies like DeepSeek-V3’s DualPipe and to compose ZeRO levels with PP/EP/DP without bespoke code.
- Tools/workflows: “Schedule templates” library (e.g., DualPipe, 1F1B, ZeroBubble), CI checks to enforce schedule correctness, migration scripts from Megatron/DeepSpeed configurations to Piper schedules.
- Assumptions/dependencies: PyTorch models must be traceable by TorchDynamo; CUDA/NCCL stack; consistent collective ordering; hardware topology awareness; network bandwidth sufficient to benefit from overlap.
- Cloud providers and Managed ML services: Strategy-agnostic training service
- What: Offer customers a declarative schedule interface and compile to per-device execution plans; support heterogeneous deployments across clusters with the same IR.
- Tools/workflows: Ray-based distributed runtime, service-level “schedule registry,” cost/throughput presets (e.g., “max-throughput,” “max-memory-efficiency”).
- Assumptions/dependencies: Kubernetes/Slurm integration; per-tenant communicator isolation; telemetry for scheduling decisions; standardized NCCL versions across fleet.
- MLOps/FinOps: Cost and energy reduction via joint compute–communication scheduling and larger effective batch sizes
- What: Reduce time-to-train and energy/CO₂ per epoch by overlapping comm/compute, optimizing bucket sizes and streams, and enabling ZeRO compositions that allow bigger batches.
- Tools/workflows: Profile-guided bucket sizing; dashboards that visualize Piper’s training DAG and stream utilization; automated “schedule sanity checks.”
- Assumptions/dependencies: Stable runtime profiling; predictable network behavior; datasets/models that benefit from larger batches.
- Academia (ML systems research): Fast prototyping of new parallelism and PP schedules
- What: Express novel intra-/inter-device schedules (e.g., new microbatch overlaps, hybrid EP/TP/CP) with a few directives; run controlled ablations.
- Tools/workflows: IR visualizer; side-by-side schedule comparison; reproducible “schedule descriptors” stored with experiments.
- Assumptions/dependencies: Access to multi-GPU nodes; students familiar with PyTorch; reliable tracing for research models.
- Academia/Education: Teaching distributed training concepts with a programmable IR
- What: Use Piper’s annotations/directives to demonstrate DP/PP/EP and schedule trade-offs (overlap vs. bubbles, stream interference).
- Tools/workflows: Classroom labs with small clusters; canned exercises that toggle directives and visualize runtime changes.
- Assumptions/dependencies: GPU access for students; simplified example models that are statically traceable.
- Healthcare/Life Sciences (domain LLMs/MoE): Train larger or more memory-efficient models on constrained budgets
- What: Compose ZeRO levels with PP to fit bigger models or batches, improving training stability and convergence for medical/biotech models.
- Tools/workflows: Predefined schedules for MoE layers vs. non-expert blocks; compliance-aligned on-prem clusters running Piper runtime.
- Assumptions/dependencies: Data governance constraints; validated NCCL/CUDA versions; model code traceability.
- HPC centers: Reproducible schedule specification and cross-hardware benchmarking
- What: Use a unified DAG IR to standardize reporting of training schedules, enabling apples-to-apples comparisons and persistence across hardware upgrades.
- Tools/workflows: “Schedule as artifact” in benchmark suites; Slurm/Ray adapters for multi-node deployment; automatic device-mesh mapping.
- Assumptions/dependencies: Accurate topology discovery; strict communicator ordering to avoid deadlocks.
- Open-source/SMEs: Port existing single-GPU or DP-only models to multi-GPU with minimal rewrites
- What: Annotate model components and apply schedule templates to scale beyond DP, benefiting projects without in-house systems teams.
- Tools/workflows: Starter kits per architecture (e.g., GPT, ViT, MoE); community repository with shared schedules.
- Assumptions/dependencies: Static graphs (limited dynamic control flow support today); basic multi-GPU hardware.
- Energy/Sustainability teams: Immediate tracking of throughput and memory-efficiency gains attributable to schedule changes
- What: Attribute energy savings to schedule edits (e.g., using same-stream DP/EP comms vs. separate streams); report improved performance-per-watt.
- Tools/workflows: Energy and utilization monitoring; Piper DAG annotations linked to metering data.
- Assumptions/dependencies: Power telemetry available; repeatable workloads for before/after comparisons.
- Governance/Reproducibility (cross-sector): Include a standardized “schedule descriptor” in publications and model cards
- What: Record placement and order directives for reproducibility and auditing; aids peer review and replication.
- Tools/workflows: Export/import of schedule directives and IR; integration with artifact repositories.
- Assumptions/dependencies: Community willingness to publish schedules; alignment with existing reproducibility checklists.
Long-Term Applications
These use cases require additional research, scaling, or development, but are directly enabled or suggested by Piper’s design (IR, directives, runtime).
- Automated schedule synthesis and search (AutoPiper)
- What: Use profile-guided optimization or learned cost models to explore directive spaces (streams, bucket sizes, overlaps) and auto-generate near-optimal schedules.
- Tools/products: Auto-tuner integrated with the compiler; schedule “blueprints” seeded by templates.
- Dependencies/assumptions: Accurate performance models; low-overhead profiling; search scalability.
- Generalized sharding inference (GSPMD-like) across directives
- What: Infer communication for arbitrary Replicate/Shard compositions, automatically propagate placements, and optimize communication elision.
- Tools/products: Advanced compiler passes; constraint solvers for placement.
- Dependencies/assumptions: Robust correctness proofs; interoperability with heterogeneous sharding (EP/TP/CP).
- Fine-grained control of intra-GPU resources
- What: Integrate SM partitioning, kernel priorities, and QoS to manage compute vs. comm interference beyond stream-level heuristics.
- Tools/products: CUDA MPS/stream priority integrations; vendor APIs for SM budgeting.
- Dependencies/assumptions: Vendor support (NVIDIA, AMD); predictable kernel interference models.
- Support for dynamic models and control flow
- What: Extend beyond fully traceable graphs to models with dynamic routing, variable sequence lengths, or conditional branches.
- Tools/products: Hybrid eager+compiled execution with dynamic IR patches; runtime guards.
- Dependencies/assumptions: Advances in PyTorch/TorchDynamo; runtime overhead constraints.
- Multi-hardware backends (e.g., TPU, AMD ROCm, custom accelerators)
- What: Retain the same IR and directives but lower to different collective libraries and device runtimes.
- Tools/products: Backend adapters for collective ops; device-mesh abstractions per platform.
- Dependencies/assumptions: Mature collective libraries (e.g., RCCL, XLA collectives); cross-platform stream/event semantics.
- Cluster-aware and network-aware scheduling
- What: Jointly plan communication and compute considering cluster topology, congestion, and co-tenancy; integrate with Kubernetes/Slurm schedulers.
- Tools/products: Topology-aware device mesh mapper; admission control guided by DAG costs.
- Dependencies/assumptions: Real-time telemetry; APIs from cluster schedulers; consistent network fabrics.
- Energy-/carbon-aware training
- What: Optimize schedules for energy efficiency or carbon intensity (time-of-day pricing, renewable availability) while meeting performance SLAs.
- Tools/products: Energy-cost models integrated with the scheduler; “green schedule” presets.
- Dependencies/assumptions: Accurate power/CO₂ signals; policy compliance for scheduling delays.
- Standardization of a Training DAG IR as an interchange format
- What: Cross-framework standard adopted by academia/industry for sharing training schedules and placements.
- Tools/products: Open specification; converters for JAX/XLA, Megatron, DeepSpeed, PyTorch.
- Dependencies/assumptions: Community alignment; governance by standard bodies or consortia.
- Security and compliance features
- What: Deterministic schedule replay, communication logs, and provenance for regulated sectors (healthcare, finance, government).
- Tools/products: Audit trail generation; signed “schedule manifests”; policy-aware directive validators.
- Dependencies/assumptions: Storage and privacy policies; deterministic seeds for randomness.
- Developer experience and debugging ecosystem
- What: IDE plugins for directive authoring, linting, and visualization; live DAG inspection and stream timelines.
- Tools/products: VS Code/JetBrains extensions; web-based DAG explorer; diff tools for schedules.
- Dependencies/assumptions: Stable IR APIs; integration with profiling tools (Nsight, PyTorch Profiler).
- Training-as-DSL for curricula and MOOCs at scale
- What: High-level “training DSL” derived from Piper directives to teach distributed systems concepts; emulators for large schedules on small clusters.
- Tools/products: Classroom simulators; auto-graders for schedule correctness.
- Dependencies/assumptions: Lightweight simulators; simplified performance models.
Each long-term item builds on Piper’s core idea—decoupling the “what” (strategy) from the “how” (runtime)—and extends it toward automation, broader hardware support, stronger guarantees, and richer ecosystem tooling.
Glossary
- All-gather: A collective communication operation that gathers distributed shards into a full tensor on all ranks. "When shard_params=True, insert an all-gather Comm before every matched node."
- All-reduce: A collective that aggregates tensors (e.g., gradients) across ranks and returns the result to all ranks. "averages the gradients across all replicas using an allreduce (or allgather and reduce-scatter)."
- All-to-all (A2A): A collective where each rank sends distinct data to every other rank, commonly used for sharded expert activations. "Inserts an all-to-all Comm before and after each matched Chunk"
- Autograd: Automatic differentiation engine that records and executes backward graphs. "an opaque graph inserted by the PyTorch autograd engine"
- Bucket size (communication): The maximum size used to group parameters/gradients into communication buckets to balance bandwidth utilization and overlap. "choosing a communication bucket size that is large enough to ensure good network utilization but small enough to enable overlap- ping with compute"
- Collective communication: Multi-rank communication primitives (e.g., all-reduce, all-gather) executed by a group of processes. "Workers execute collective communication according to the sharding plan"
- Communicator: The per-rank context used to coordinate ordering and membership of collective operations. "GPU collective communication typically requires a 'communicator'"
- Context parallelism (CP): Parallelism that partitions the input sequence (context) across devices. "context (CP) and pipeline (PP) parallelism"
- CUDA stream: A sequence of GPU operations that execute in order; different streams can overlap. "scheduler for operators over intra-GPU resources such as CUDA streams."
- Data parallelism (DP): Replicates model weights across devices and averages gradients across replicas. "In DP [3, 22], each worker holds a replica of the model weights"
- Device mesh: An N-dimensional arrangement of device identifiers used for placement and sharding specifications. "takes in a device mesh (an N-D array of devices)"
- Directed acyclic graph (DAG): A graph with directed edges and no cycles; used here to represent training compute and communication dependencies. "a unified global training DAG that represents all computation and communication"
- DualPipe: A pipeline-parallel schedule that overlaps forward and backward microbatches to hide communication, especially when composed with EP. "DeepSeek- V3 introduced DualPipe, a custom PP schedule"
- DualPipeV: A variant of DualPipe placing each pipeline stage on one PP rank while retaining microbatch overlap. "This strategy is a variant of DualPipe known as DualPipeV [35]"
- Expert parallelism (EP): Shards model experts across devices and routes tokens to a subset of experts. "a subset of experts in EP [21]"
- Fully sharded DP (FSDP): A form of data parallelism that shards parameters, gradients, and optimizer state across ranks. "also known as fully sharded DP [55]"
- GSPMD: A tensor partitioning and placement framework for composing sharding/replication strategies. "GSPMD [49]"
- High Bandwidth Memory (HBM): On-package memory with very high throughput used by GPUs. "SMs and HBM or network bandwidth"
- Mixture-of-Experts (MoE): A model architecture where inputs are routed to a subset of expert networks to scale capacity efficiently. "applied to an MoE transformer model"
- MPMD (Multiple Program Multiple Data): A distributed model where different ranks run different programs, often used in pipeline parallelism. "most easily executed with an MPMD (multiple program multiple data) approach"
- NCCL: NVIDIA’s library for efficient collective communication across GPUs. "or a NCCL communication kernel inserted by Piper."
- Pipeline parallelism (PP): Splits model layers across devices and pipelines microbatches through stages. "PP schedules use multiple microbatches to overlap execution across devices."
- Point-to-point (P2P): Direct send/receive communication between a pair of devices. "add a P2P send/recv Comm"
- Ray: A distributed execution framework used here to implement Piper’s runtime. "use Ray to implement the distributed runtime."
- Reduce-scatter: A collective that reduces inputs across ranks and scatters the result shards, often used for gradient sharding. "reduce-scatter in the case of gradients."
- Rematerialization: Temporarily reconstructing full tensors (e.g., parameters) from shards when needed for computation. "rematerial- ized with allgather when needed"
- Streaming multiprocessors (SMs): GPU compute units that execute kernels; their allocation affects overlap and performance. "streaming multiprocessors (SMs) allocated to compute vs. communication."
- SPMD (Single Program Multiple Data): A model where all ranks run the same program on different data/partitions. "SPMD-style tensor annotations"
- Tensor parallelism (TP): Splits tensors (e.g., weight matrices) across devices to parallelize large operations. "rows/columns in TP [40]"
- TorchDynamo: A PyTorch JIT compiler that captures Python and tensor operations into FX graphs. "TorchDynamo [4], a PyTorch JIT compiler"
- torch.compile: PyTorch compilation entry point used to integrate Piper as a backend. "as a torch.compile [4] backend"
- ZeRO: A family of techniques that shard optimizer state, gradients, and/or parameters to reduce memory redundancy. "ZeRO [38], also known as fully sharded DP"
- ZeRO-2: The ZeRO variant that shards gradients (and optionally optimizer state), often called “ZeRO stage 2.” "ZeRO-2 gradient sharding"
- ZeroBubble: A pipeline schedule that splits backward passes to reduce pipeline bubbles. "ZeroBubble [34], a PP schedule"
Collections
Sign up for free to add this paper to one or more collections.
