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.

Aller à l'installation

Dépôt: hiyenwong/ai_collection
Dernière activité de la source: 7 juillet 2026 à 08:26
Langue détectée de SKILL.md: anglais
Étoiles: 2
Forks: 0

Options d'installation

Le prompt qui vérifie d'abord la source est sélectionné par défaut. Vous pouvez passer à une commande directe ou télécharger une copie locale.

Vérifiez les fichiers source

Lisez SKILL.md et les fichiers associés affichés par SkillsMP avant de décider de l'installer.

Retour à la recherche Retour à l’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