piper-programmable-distributed-training

Piper framework for user-controllable distributed training that decouples parallelism strategy from runtime implementation using unified global training DAG intermediate representation. Use for distributed ML training, parallelism strategy design, and flexible training system architecture.

Jump to install

Repository: hiyenwong/ai_collection
Last source activity: July 7, 2026 at 08:26
Detected SKILL.md language: English
Stars: 2
Forks: 0

Install options

The review-first prompt is selected by default. You can switch to a direct command or download a local copy.

Review the source files

Read SKILL.md and any companion files shown by SkillsMP before deciding whether to install.

Back to search Back to installation

name	piper-programmable-distributed-training
description	Piper framework for user-controllable distributed training that decouples parallelism strategy from runtime implementation using unified global training DAG intermediate representation. Use for distributed ML training, parallelism strategy design, and flexible training system architecture.
version	1.0.0
category	systems-engineering
authors	["Megan Frisella","Shubham Tiwari","Andy Ruan","Yi Pan","Parker Gustafson","Mat Jacob","Gilbert Bernstein","Stephanie Wang"]
arxiv_id	2606.11169
submission_date	2026-06-09
activation_keywords	["Piper","distributed training","parallelism strategy","ZeRO","pipeline parallelism","expert parallelism","training DAG","DualPipe","DeepSeek-V3","strategy-implementation decoupling"]

Piper: Programmable Distributed Training System

Core Contribution

Fundamental Problem: Existing systems require manual strategy design + implementation, making adaptation difficult. General-purpose frameworks are tied to fixed parallelism strategies, hindering state-of-the-art integration.

Solution: Decouple strategy from runtime implementation using:

User declarations: Model annotations + scheduling directives
Unified IR: Global training DAG representing all computation/communication
Strategy-agnostic runtime: Executes compiled per-device plans

Architecture

1. Three-Level Abstraction

┌─────────────────────────────────────────┐
│  User Strategy Declaration              │  ← High-level annotations
│  (model annotations + directives)        │
└─────────────────────────────────────────┘
                    ↓
┌─────────────────────────────────────────┐
│  Piper Intermediate Representation (IR) │  ← Unified global training DAG
│  (compute + communication operations)    │
└─────────────────────────────────────────┘
                    ↓
┌─────────────────────────────────────────┐
│  Per-Device Execution Plans             │  ← Compiled device-specific schedules
│  (strategy-agnostic runtime)             │
└─────────────────────────────────────────┘

2. Key Components

A. Strategy Declaration Layer

# Example: DualPipe strategy declaration (DeepSeek-V3 style)
@piper.annotate(
    data_parallel=4,
    pipeline_parallel=8,
    expert_parallel=2,
    zero_stage=3
)
@piper.directive("DualPipe", schedule="interleaved_forward_backward")
class MyModel:
    def forward(self, x):
        # Model logic
        ...

B. Intermediate Representation (IR)

class PiperIR:
    
    
     ():
        .dag = ._build_global_dag(model_annotations)
        ._apply_directives(directives)
    
     ():
        
        ...
    
     ():
        
         directive  directives:
            .dag = directive.transform(.dag)

def compile_per_device_plans(ir, num_devices):
    """
    Generate device-specific execution schedules.
    
    Inputs:
    - ir: Global training DAG
    - num_devices: Total devices
    
    Outputs:
    - device_plans[device_id]: Execution schedule for each device
    """
    device_plans = {}
    
    for device_id in range(num_devices):
        # Extract device-relevant operations
        device_ops = extract_device_operations(ir, device_id)
        
        # Schedule operations respecting dependencies
        schedule = schedule_operations(device_ops)
        
        device_plans[device_id] = schedule
    
    return device_plans

class PiperRuntime:
    """
    Distributed runtime that executes compiled plans.
    Independent of strategy - only executes scheduled operations.
    """
    
    def execute_plan(self, device_plan):
        for operation in device_plan:
            if operation.type == "compute":
                self.execute_compute(operation)
            elif operation.type == "communication":
                self.execute_communication(operation)

@piper.annotate(data_parallel=N, zero_stage=3)
@piper.directive("ZeRO", partition="optimizer_states+gradients+parameters")

@piper.annotate(pipeline_parallel=P)
@piper.directive("Pipeline", schedule="GPipe")  # or "1F1B"

@piper.annotate(expert_parallel=E)
@piper.directive("ExpertParallel", all_to_all="router+output")

@piper.annotate(
    data_parallel=4,
    pipeline_parallel=8,
    expert_parallel=2
)
@piper.directive("DualPipe", schedule="interleaved_forward_backward")

import piper

@piper.annotate(
    data_parallel=4,        # Number of data parallel replicas
    pipeline_parallel=8,    # Number of pipeline stages
    expert_parallel=2,      # Number of expert parallel groups
    zero_stage=3            # ZeRO optimization stage
)
class TransformerModel:
    def __init__(self, config):
        self.layers = [TransformerLayer(config) for _ in range(config.num_layers)]
    
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

# Directive: transforms IR based on strategy
@piper.directive("DualPipe", 
    schedule="interleaved_forward_backward",
    communication_overlap=True,
    memory_optimization="activation_checkpointing"
)

# Piper automatically generates unified DAG
ir = piper.generate_ir(model_annotations, directives)

# DAG structure
ir.nodes = [
    ComputeOp("forward_layer_0", device=0),
    ComputeOp("forward_layer_1", device=1),
    CommunicationOp("send_activation_0_to_1", src=0, dst=1),
    ComputeOp("backward_layer_1", device=1),
    CommunicationOp("all_reduce_gradients", devices=[0,1,2,3]),
    ...
]

ir.edges = [
    Dependency("forward_layer_0", "send_activation_0_to_1"),
    Dependency("send_activation_0_to_1", "forward_layer_1"),
    ...
]

# Compile per-device execution plans
device_plans = piper.compile(ir, num_devices=32)

# Example device plan
device_plans[0] = [
    {"op": "forward_layer_0", "time": 0},
    {"op": "send_activation", "time": 1},
    {"op": "backward_layer_0", "time": 2},
    {"op": "all_reduce_gradients", "time": 3},
    {"op": "optimizer_step", "time": 4}
]

# Strategy-agnostic execution
runtime = PiperRuntime(num_devices=32)
runtime.execute(device_plans)

Strategy	Traditional Framework	Piper	Improvement
ZeRO-3	1.0x baseline	1.0x	Parity (maintained)
Pipeline	Manual tuning required	Auto-optimized	Faster deployment
DualPipe	Not supported	1.15x throughput	Novel strategy enabled

Aspect	Traditional	Piper
Strategy implementation	Hand-coded runtime	declarative directives
Strategy adaptation	Code refactoring	Annotation modification
Novel strategies	Custom runtime needed	Directive + IR transform
Compute-communication scheduling	Sequential	Joint optimization
Memory optimization	Manual tuning	IR-level scheduling

@piper.annotate(data_parallel=128, pipeline_parallel=8, zero_stage=3)
@piper.directive("ZeRO-3", communication_overlap=True)
# Large-scale GPT-style training

@piper.annotate(expert_parallel=16, data_parallel=4)
@piper.directive("ExpertParallel", load_balancing="dynamic")
# Mixture-of-experts training

@piper.annotate(data_parallel=4, pipeline_parallel=8, expert_parallel=2)
@piper.directive("DualPipe", interleaved_schedule=True)
# Composed parallelism with communication overlap

# Compare Piper ZeRO-3 vs traditional ZeRO-3
traditional_time = train_with_traditional_zero3(model, data)
piper_time = train_with_piper(model, data, annotations={"zero_stage": 3})

assert abs(piper_time - traditional_time) < tolerance  # Performance parity