Pipeline Parallelism

import torch
from torch import nn
from d9d.pipelining.api import PipelineStageInfo, distribute_layers_for_pipeline_stage

class MyModelChunk(nn.Module):
    def __init__(self, stage: PipelineStageInfo, config):
        super().__init__()
        self.stage = stage
        self.config = config

        # 1. Determine what layers live here
        self.start_layer, self.end_layer = distribute_layers_for_pipeline_stage(
            config.n_layers, num_virtual_layers_pre=1, num_virtual_layers_post=1, stage=stage
        )

        # 2. Build sub-modules (using ModuleDict - for compatibility)
        self.layers = nn.ModuleDict({
            str(layer): TransformerBlock(...) 
            for layer in range(self.start_layer, self.end_layer)
        })

        # Only build embeddings on first stage
        if stage.is_current_stage_first:
            self.embed = nn.Embedding(...)

        # Only build head on last stage
        if stage.is_current_stage_last:
            self.head = nn.Linear(...)

    def forward(self, input_ids=None, hidden_states=None):        
        # Run embeddings only on first stage
        if self.stage.is_current_stage_first:
            x = self.embed(input_ids)
        else:
            x = hidden_states

        # Run local layers
        for layer_idx in range(self.start_layer, self.end_layer):
            x = self.layers[str(layer_idx)](x)

        outputs = {
            "hidden_states": x
        }

        # Last stage logic
        if self.stage.is_current_stage_last:
            logits = self.head(x)
            outputs['logits'] = logits

        return outputs

    # --- Protocol Implementation ---

    def infer_stage_inputs_from_pipeline_inputs(self, inputs: dict[str, torch.Tensor], n_microbatches: int):
        batch_size = inputs['input_ids'].shape[0]
        micro_batch_size = batch_size // n_microbatches
        seq_len = inputs['input_ids'].shape[1]

        if self.stage.is_current_stage_first:
            # First stage receives raw input IDs
            return {"input_ids": torch.empty((micro_batch_size, seq_len), dtype=torch.long)}
        else:
            # Intermediate stages receive hidden states from previous stage
            return {"hidden_states": torch.empty((micro_batch_size, seq_len, self.hidden_dim))}

    def infer_stage_outputs_from_pipeline_inputs(self, inputs: dict[str, torch.Tensor], n_microbatches: int):
        batch_size = inputs['input_ids'].shape[0]
        micro_batch_size = batch_size // n_microbatches
        seq_len = inputs['input_ids'].shape[1]

        outputs = {"hidden_states": torch.empty((micro_batch_size, seq_len, self.config.hidden_dim))}

        if self.stage.is_current_stage_last:
            # Last stage outputs logits too
            outputs["logits"] = torch.empty((micro_batch_size, seq_len, self.config.vocab_size))

        return outputs

from d9d.pipelining.api import PipelineShardingSpec
from d9d.core.sharding import ShardingSpec
from torch.distributed.tensor import Shard

# Example: Split 'images' on dim 1, but replicate 'camera_angles' across all microbatches
my_spec = PipelineShardingSpec(
    input_data={
        "images": Shard(1),
        "camera_angles": None
    }
    # input_kwargs can be defined similarly
)

from torch import Tensor
from torch.distributed.tensor import Shard
import torch.nn.functional as F

from d9d.core.dist_context import DistributedContext
from d9d.core.sharding import shard_tree
from d9d.pipelining.factory import build_schedule, PipelineSchedule1F1BConfig
from d9d.pipelining.api import PipelineShardingSpec


# 0. Define an object that manages loss calculation across steps
class PipelineLossHandler:
    def __init__(self, num_microbatches: int):
        self._shard_spec = {
            'target': Shard(0)
        }
        self._num_microbatches = num_microbatches
        self._targets = None

    def set_targets(self, targets: Tensor):
        self._targets = shard_tree(
            {'target': targets},
            sharding_spec=self._shard_spec,
            num_shards=self._num_microbatches,
            enforce_even_split=True
        )

    def compute_loss(self, outputs: dict[str, Tensor], microbatch_idx: int):
        # Implement any custom logic here
        current_target = self._targets[microbatch_idx]
        return F.cross_entropy(outputs['logits'].view(-1, outputs['logits'].shape[-1]), current_target.view(-1))


# 1. Define configuration
dist_context: DistributedContext = ...
model_config = ...
n_microbatches = 32
schedule_config = PipelineSchedule1F1BConfig(
    num_stages_per_rank=4,  # 4 Virtual stages per rank
    zero_bubble=True  # Enable ZB1P optimization
)

# 2. Build the schedule, model shards and loss compute function
loss_handler = PipelineLossHandler(num_microbatches=n_microbatches)
schedule_info, modules = build_schedule(
    dist_context=dist_context,
    n_microbatches=32,
    schedule_config=schedule_config,
    model_provider=lambda stage: MyModelChunk(stage, model_config),  # Factory function
    callback=loss_handler.compute_loss
)

# 3. Execution
# The schedule object exposes a simple step API
inputs = {"input_ids": ...}  # Full batch
loss_handler.set_targets(...)  # Set targets for full batch
schedule_info.schedule.configure_buffers(  # Pre-allocate buffers
    inputs, 
    kwargs={}, 
    sharding_spec=PipelineShardingSpec()
)
schedule_info.schedule.step(inputs, kwargs={})