HyperParallel Graph Mode

Graph-mode architecture for automatic parallelization with FSDP.

Note: The compile/ subpackage relies on a patched autograd engine for joint-graph capture (see Limitations). Its torch.* imports — including torch.fx.experimental and torch._guards — are therefore intentional and carry a forbidden-backend-import suppression in .jenkins/check/config/filter_pylint.txt.

Core Concept

User: Write model code + parallel configuration Framework: Graph capture → FSDP partitioning → Communication-compute overlap → Execution

Architecture

Layer 1: Configuration

from hyper_parallel.compile import GraphParallelPlan, PassConfig

# Configure which modules to mark for FSDP
parallel_plan = GraphParallelPlan()
parallel_plan.fsdp_mark("tok_embeddings")
parallel_plan.fsdp_mark_pattern("layers.*")

# Parallel configuration
pass_config = PassConfig(enable_overlap=True)

Layer 2: Graph Capture

trace_model_graph captures forward + backward into a joint FX graph:

• Parameters/buffers are static inputs (placeholders), not get_attr nodes
• torch.autograd.grad runs inside the traced function
• Uses FakeTensorMode + make_fx for symbolic tracing

Layer 3: Pass Pipeline

DeadCodeElimination → CanonicalizeGraph → FSDPPass → AutoOverlapPass

FSDPPass:

• Identifies FSDP parameter placeholders via GraphParallelPlan
• Inserts all_gather after each FSDP parameter (Shard → Replicate)
• Inserts reduce_scatter on gradient outputs (Replicate → Shard)
• Physically shards live model parameters (dim 0) so optimizer is FSDP-agnostic

AutoOverlapPass:

• Reorders wait_tensor nodes for communication-compute overlap

Layer 4: Execution

from hyper_parallel.compile import GraphTrainer

trainer = GraphTrainer(
    model=model,
    train_fn=train_fn,
    pass_config=pass_config,
    parallel_plan=parallel_plan,
)

# Compile on first batch, then run forward + backward + optimizer
trainer.train(dataloader, max_steps=100, log_interval=10)

The compiled graph is a GraphModule that:

• Takes sharded parameters as inputs
• Gathers them via AllGather at each step
• Computes forward + backward
• Scatters gradients via ReduceScatter
• Outputs loss + grads

Usage

from hyper_parallel.compile import GraphParallelPlan, GraphTrainer, PassConfig

# 1. Model
model = Llama3Model(config)

# 2. Parallel plan
parallel_plan = GraphParallelPlan()
parallel_plan.fsdp_mark_pattern("layers.*")

# 3. Trainer
trainer = GraphTrainer(
    model=model,
    train_fn=lambda m, x, y: m(x).loss(y),
    pass_config=PassConfig(enable_overlap=True),
    parallel_plan=parallel_plan,
)

# 4. Training
trainer.train(dataloader, max_steps=1000)

Key Design Decisions

1. Static Inputs: Parameters are graph inputs, not get_attr. This allows passes to split the graph by reshaping placeholders.

2. Joint Graph: Forward + backward captured together via torch.autograd.grad inside the traced function.

3. FSDP-Agnostic Optimizer: FSDPPass shards the live model's parameters in place. model.parameters() returns shards, so optimizer needs no FSDP awareness.

4. Declarative Sharding: GraphParallelPlan uses FQN patterns (layers.*) instead of imperative module wrapping.

Limitations

• FSDP only (TP/EP/PP planned)
• Parameters must have dim 0 divisible by world_size
• Requires torch with patched autograd engine for joint-graph capture