Pipeline Parallelism with AutoPipeline

1
import torch
2
from torch.distributed.device_mesh import init_device_mesh
3
from nemo_automodel.components.distributed.pipelining import AutoPipeline
4
from nemo_automodel.components.utils.model_utils import init_empty_weights
5
from transformers import AutoModelForCausalLM
6
from transformers.initialization import no_init_weights
7
from transformers.utils import ContextManagers
8
9
def loss_fn(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
10
    """Define loss function for pipeline training."""
11
    return torch.nn.functional.cross_entropy(
12
        logits.float().view(-1, logits.size(-1)),
13
        targets.view(-1),
14
        ignore_index=-100
15
    )
16
17
if __name__ == "__main__":
18
    # 1) Initialize device mesh with 2 pipeline stages
19
    world_mesh = init_device_mesh("cuda", mesh_shape=(2,), mesh_dim_names=("pp",))
20
21
    # 2) Load model on meta device to avoid OOM with large models
22
    init_ctx = ContextManagers([no_init_weights(), init_empty_weights()])
23
    with init_ctx:
24
        model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")
25
26
    # 3) Configure and build pipeline
27
    ap = AutoPipeline(
28
        world_mesh=world_mesh,
29
        pp_axis_name="pp",
30
        pp_schedule="1f1b",
31
        pp_microbatch_size=1,
32
        pp_batch_size=8,  # Local batch size consumed by the pipeline schedule
33
        device=torch.cuda.current_device(),
34
        dtype=torch.bfloat16,
35
    ).build(model, loss_fn=loss_fn)
36
37
    # 4) Access pipeline components
38
    print(ap.debug_summary())
39
    print(ap.pretty_print_stages())

1
ap = AutoPipeline(
2
    # Device mesh configuration
3
    world_mesh=world_mesh,           # DeviceMesh with pipeline axis
4
    pp_axis_name="pp",              # Name of pipeline axis (default: "pp")
5
6
    # Schedule configuration
7
    pp_schedule="1f1b",             # Pipeline schedule ("1f1b", "LoopedBFS", etc.)
8
    pp_microbatch_size=1,           # Microbatch size per stage
9
    pp_batch_size=8,                # Local schedule batch, split into microbatches
10
    # Recipe flows set this from step_scheduler.local_batch_size
11
12
    # Stage configuration
13
    layers_per_stage=None,          # Layers per stage (None for auto)
14
    module_fqns_per_model_part=None,  # Manual module assignment
15
16
    # Model patching (HF-specific)
17
    patch_inner_model=True,         # Make decoder forward stage-friendly
18
    patch_causal_lm_model=True,     # Make CausalLM wrapper return tensors (hidden/logits)
19
).build(model, loss_fn=loss_fn)

1
ap = AutoPipeline(
2
    world_mesh=world_mesh,
3
    pp_schedule="1f1b",
4
    layers_per_stage=8,  # Each stage gets ~8 transformer layers
5
).build(model, loss_fn=loss_fn)

1
from nemo_automodel.components.distributed.pipelining.functional import (
2
    generate_hf_model_fqn_per_model_part
3
)
4
5
# Generate balanced assignments
6
module_fqns = generate_hf_model_fqn_per_model_part(
7
    num_stages=4,
8
    num_layers=32,
9
    include_embeddings=True,
10
    include_lm_head=True,
11
    include_rotary_emb=True,
12
    fqn_prefix="model."
13
)
14
15
# Or define custom assignments
16
custom_module_fqns = [
17
    # Stage 0: Embeddings + first 8 layers
18
    ["model.embed_tokens", "model.rotary_emb"] +
19
    [f"model.layers.{i}" for i in range(8)],
20
21
    # Stage 1: Next 8 layers
22
    ["model.rotary_emb"] + [f"model.layers.{i}" for i in range(8, 16)],
23
24
    # Stage 2: Next 8 layers
25
    ["model.rotary_emb"] + [f"model.layers.{i}" for i in range(16, 24)],
26
27
    # Stage 3: Final 8 layers + output
28
    ["model.rotary_emb"] + [f"model.layers.{i}" for i in range(24, 32)] +
29
    ["model.norm", "lm_head"]
30
]
31
32
ap = AutoPipeline(
33
    world_mesh=world_mesh,
34
    module_fqns_per_model_part=custom_module_fqns,
35
).build(model, loss_fn=loss_fn)

1
# Stage 0 (Rank 0): Input processing + first half
2
stage_0_modules = [
3
    "model.embed_tokens",     # Token embeddings
4
    "model.layers.0-15",      # First 16 transformer layers
5
    "model.rotary_emb"        # Position embeddings (shared)
6
]
7
8
# Stage 1 (Rank 1): Second half + output processing
9
stage_1_modules = [
10
    "model.layers.16-31",     # Last 16 transformer layers
11
    "model.norm",             # Final layer norm
12
    "lm_head",               # Language modeling head
13
    "model.rotary_emb"        # Position embeddings (shared)
14
]

1
# Stage 0 (Rank 0): Input processing
2
stage_0_modules = [
3
    "model.embed_tokens",     # Token embeddings
4
    "model.layers.0-7",       # First 8 transformer layers
5
    "model.rotary_emb"        # Position embeddings (shared)
6
]
7
8
# Stage 1 (Rank 1): Early layers
9
stage_1_modules = [
10
    "model.layers.8-15",      # Next 8 transformer layers
11
    "model.rotary_emb"
12
]
13
14
# Stage 2 (Rank 2): Middle layers
15
stage_2_modules = [
16
    "model.layers.16-23",     # Next 8 transformer layers
17
    "model.rotary_emb"
18
]
19
20
# Stage 3 (Rank 3): Output processing
21
stage_3_modules = [
22
    "model.layers.24-31",     # Final 8 transformer layers
23
    "model.norm",             # Final layer norm
24
    "lm_head",               # Language modeling head
25
    "model.rotary_emb"
26
]

1
from nemo_automodel.components.distributed.pipelining.functional import stage_ids_this_rank
2
3
# Calculate which stages run on this rank
4
# For a "loop" style schedule (default)
5
stage_ids = stage_ids_this_rank(pp_rank=0, pp_size=4, num_stages=8, style="loop")
6
# Returns: (0, 4) - rank 0 gets stages 0 and 4
7
8
# For a "v" style schedule (for zero-bubble schedules)
9
stage_ids = stage_ids_this_rank(pp_rank=0, pp_size=4, num_stages=8, style="v")
10
# Returns: (0, 7) - rank 0 gets stages 0 and 7

1
from nemo_automodel.components.distributed.pipelining.functional import (
2
    generate_hf_model_fqn_per_model_part
3
)
4
5
# Generate balanced module assignments for an HF-style decoder layout
6
module_names = generate_hf_model_fqn_per_model_part(
7
    num_stages=4,
8
    num_layers=32,
9
    include_embeddings=True,
10
    include_lm_head=True,
11
    include_rotary_emb=False,  # Set based on your model
12
    fqn_prefix=""  # Use "model." for nested models
13
)

1
from nemo_automodel.components.distributed.pipelining.functional import calculate_virtual_stages
2
3
# Calculate virtual stages for interleaved schedules
4
num_virtual_stages, stages_per_rank = calculate_virtual_stages(
5
    num_layers=32,
6
    layers_per_stage=4,  # Each virtual stage has 4 layers
7
    pp_size=4,
8
    is_single_stage_schedule=False,
9
    round_to_pp_multiple="up"  # Round up to nearest multiple of pp_size
10
)

1
from nemo_automodel.components.distributed.pipelining.functional import build_pipeline_schedule
2
3
# Build a schedule for your stages
4
schedule = build_pipeline_schedule(
5
    pipeline_parallel_schedule_csv=None,  # Optional CSV schedule
6
    pipeline_parallel_schedule="1f1b",
7
    microbatch_size=1,
8
    local_batch_size=8,
9
    stages=stages,  # List of PipelineStage objects
10
    loss_fn=loss_fn,
11
    scale_grads=False
12
)

1
# NON-RUNNABLE PSEUDOCODE
2
import torch
3
import torch.nn as nn
4
from torch.distributed.device_mesh import init_device_mesh
5
from torch.distributed.pipelining import PipelineStage
6
from nemo_automodel.components.distributed.pipelining.functional import (
7
    stage_ids_this_rank,
8
    build_pipeline_schedule,
9
    calculate_virtual_stages
10
)
11
12
class CustomTransformerBlock(nn.Module):
13
    def __init__(self, hidden_size):
14
        super().__init__()
15
        self.attention = nn.MultiheadAttention(hidden_size, num_heads=8)
16
        self.mlp = nn.Sequential(
17
            nn.Linear(hidden_size, hidden_size * 4),
18
            nn.GELU(),
19
            nn.Linear(hidden_size * 4, hidden_size)
20
        )
21
        self.norm1 = nn.LayerNorm(hidden_size)
22
        self.norm2 = nn.LayerNorm(hidden_size)
23
24
    def forward(self, x):
25
        # Simplified transformer block
26
        attn_out, _ = self.attention(x, x, x)
27
        x = self.norm1(x + attn_out)
28
        x = self.norm2(x + self.mlp(x))
29
        return x
30
31
class CustomModel(nn.Module):
32
    def __init__(self, vocab_size, hidden_size, num_layers):
33
        super().__init__()
34
        self.embedding = nn.Embedding(vocab_size, hidden_size)
35
        self.layers = nn.ModuleList([
36
            CustomTransformerBlock(hidden_size) for _ in range(num_layers)
37
        ])
38
        self.output_proj = nn.Linear(hidden_size, vocab_size)
39
40
    def forward(self, input_ids):
41
        x = self.embedding(input_ids)
42
        for layer in self.layers:
43
            x = layer(x)
44
        return self.output_proj(x)
45
46
def split_custom_model_for_pipeline(model, pp_rank, pp_size, num_stages):
47
    """Split a custom model into pipeline stages."""
48
49
    # Determine which stages this rank handles
50
    stage_indices = stage_ids_this_rank(pp_rank, pp_size, num_stages, style="loop")
51
52
    stages = []
53
    for stage_idx in stage_indices:
54
        # Create a stage-specific version of the model
55
        # This is a simplified example - you'd need to implement proper splitting
56
        stage_model = create_stage_model(model, stage_idx, num_stages)
57
58
        # Create PipelineStage
59
        stage = PipelineStage(
60
            stage_model,
61
            stage_idx,
62
            num_stages,
63
            device=torch.cuda.current_device(),
64
            group=None  # Set your process group here
65
        )
66
        stages.append(stage)
67
68
    return stages
69
70
# Sketch of model-owned orchestration
71
def main():
72
    # Initialize device mesh
73
    world_mesh = init_device_mesh("cuda", mesh_shape=(4,), mesh_dim_names=("pp",))
74
    pp_rank = world_mesh["pp"].get_local_rank()
75
    pp_size = world_mesh["pp"].size()
76
77
    # Create model
78
    model = CustomModel(vocab_size=50000, hidden_size=768, num_layers=24)
79
80
    # Calculate virtual stages
81
    num_virtual_stages, stages_per_rank = calculate_virtual_stages(
82
        num_layers=24,
83
        layers_per_stage=3,  # 8 virtual stages total
84
        pp_size=4,
85
        is_single_stage_schedule=False
86
    )
87
88
    # Split model into stages
89
    stages = split_custom_model_for_pipeline(model, pp_rank, pp_size, num_virtual_stages)
90
91
    # Define loss function
92
    def loss_fn(logits, targets):
93
        return nn.functional.cross_entropy(
94
            logits.view(-1, logits.size(-1)),
95
            targets.view(-1)
96
        )
97
98
    # Build pipeline schedule
99
    schedule = build_pipeline_schedule(
100
        pipeline_parallel_schedule_csv=None,
101
        pipeline_parallel_schedule="interleaved1f1b",  # Multi-stage schedule registered by PyTorch
102
        microbatch_size=1,
103
        local_batch_size=8,
104
        stages=stages,
105
        loss_fn=loss_fn,
106
        scale_grads=True
107
    )
108
109
    # Training loop
110
    for batch in dataloader:
111
        # Use schedule.step() for training
112
        losses = []
113
        schedule.step(batch["input_ids"], target=batch["labels"], losses=losses)
114
115
        # losses will contain the loss values from the last stage
116
        if losses:
117
            print(f"Loss: {sum(losses) / len(losses)}")

1
from nemo_automodel.components.distributed.pipelining.functional import pipeline_model
2
3
def custom_parallelize_fn(
4
    model, world_mesh, moe_mesh, *,
5
    dp_axis_names,
6
    cp_axis_name=None,
7
    tp_axis_name=None,
8
    ep_axis_name=None,
9
    ep_shard_axis_names=None,
10
):
11
    """Custom parallelization function for each local model part."""
12
    # Apply your custom parallelization logic here
13
    # This is called for each local model part
14
    if dp_axis_names:
15
        # Apply data parallelism
16
        pass
17
    # Add any other parallelization strategies
18
    pass
19
20
# Use pipeline_model only with an HF-compatible model.
21
schedule, model_parts, has_first, has_last, stages = pipeline_model(
22
    model=hf_compatible_model,
23
    world_mesh=world_mesh,
24
    moe_mesh=None,
25
    pp_axis_name="pp",
26
    dp_axis_names=("dp",),
27
    layers_per_stage=4,
28
    pipeline_parallel_schedule="1f1b",
29
    pipeline_parallel_schedule_csv=None,
30
    microbatch_size=1,
31
    local_batch_size=8,
32
    device=torch.cuda.current_device(),
33
    loss_fn=loss_fn,
34
    parallelize_fn=custom_parallelize_fn,
35
    module_fqns_per_model_part=None,  # Or provide an explicit HF module assignment
36
    patch_inner_model=False,  # Use only when the model's inner forward is already stage-friendly
37
    patch_causal_lm_model=False,  # Use only when the wrapper already returns stage tensors
38
)

1
def parallelize_fn(
2
    model, world_mesh, moe_mesh, *,
3
    dp_axis_names,
4
    cp_axis_name=None,
5
    tp_axis_name=None,
6
    ep_axis_name=None,
7
    ep_shard_axis_names=None,
8
):
9
    """Apply additional parallelization to each local model part."""
10
    # Example: Apply FSDP to each local model part
11
    if dp_axis_names:
12
        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
13
        # Wrap model with FSDP (simplified example)
14
        # In practice, you'd configure FSDP parameters
15
        pass
16
17
    # Example: Apply tensor parallelism
18
    if tp_axis_name:
19
        # Apply tensor parallelism to attention/MLP layers
20
        pass
21
22
# Build pipeline with custom parallelization
23
ap = AutoPipeline(world_mesh=world_mesh).build(
24
    model,
25
    loss_fn=loss_fn,
26
    parallelize_fn=parallelize_fn
27
)

1
# Get pipeline info
2
info = ap.info
3
print(f"Pipeline enabled: {info.enabled}")
4
print(f"Has first stage: {info.has_first_stage}")
5
print(f"Has last stage: {info.has_last_stage}")
6
7
# Access model parts and runtime pipeline stages
8
model_parts = ap.parts  # List of local model-part nn.Module objects
9
pipeline_stages = info.stages  # List of local PipelineStage objects used by the schedule
10
stage_modules = ap.list_stage_modules()  # Module names within each local model part

1
# Parameter distribution
2
stage_param_counts = ap.get_stage_param_counts()
3
total_params = ap.get_total_param_count()
4
trainable_params = ap.get_total_param_count(trainable_only=True)
5
6
for i, params in enumerate(stage_param_counts):
7
    percentage = (params / total_params) * 100
8
    print(f"Stage {i}: {params:,} parameters ({percentage:.1f}%)")
9
10
# Debug summary
11
print(ap.debug_summary())
12
print(ap.pretty_print_stages(max_modules_per_stage=10))
13
14
# Visualize schedule
15
ap.visualize_current_schedule("pipeline_schedule.png")

1
from nemo_automodel.components.distributed.pipelining.functional import scale_grads_by_divisor
2
from nemo_automodel.components.training.utils import clip_grad_norm
3
4
# Call these after backward on every pipeline rank.
5
scale_grads_by_divisor(ap.info.stages, divisor=8)
6
7
# Use the PP-aware clipping helper across the local model parts.
8
grad_norm = clip_grad_norm(
9
    1.0,
10
    ap.parts,
11
    norm_type=2.0,
12
    pp_enabled=True,
13
    device_mesh=world_mesh,
14
    pp_axis_name="pp",
15
)