Pipeline Parallelism with AutoPipeline

Introduction

As large language models continue to grow in size, training and fine-tuning them efficiently across multiple GPUs has become increasingly challenging. While data parallelism works well for smaller models, models with billions of parameters require more sophisticated parallelization strategies to overcome memory constraints and communication overhead.

Pipeline parallelism addresses these challenges by splitting a model’s layers across different devices and processing them in a pipelined fashion. Each device processes a different stage of the model, enabling training of models that would not fit on a single device while maintaining high GPU utilization through overlapped computation.

AutoPipeline is NeMo AutoModel’s high-level pipeline parallelism interface for Hugging Face-compatible models. Built on PyTorch’s native torch.distributed.pipelining, it expects a Hugging Face-style model with a config, a supported decoder-layer layout, and pipeline-compatible forwards. Recipe-driven pipeline parallelism also requires the selected model implementation to advertise a _pp_plan (or use the supported MoE path) and pass the repository’s pipeline validation; encoder-decoder models and models with actually tied input/output embedding weights are rejected.

The functional module at nemo_automodel.components.distributed.pipelining.functional exposes the same Hugging Face-specific pipeline_model() splitter plus lower-level stage and schedule helpers. Custom PyTorch architectures can reuse helpers such as stage_ids_this_rank(), calculate_virtual_stages(), and build_pipeline_schedule() after their own code constructs valid PipelineStage objects; pipeline_model() is not a generic nn.Module splitter.

This guide walks you through AutoPipeline for compatible Hugging Face models and the lower-level helpers for manually constructed stages. You will learn how to configure pipeline stages, integrate with existing training workflows, optimize performance, and combine pipeline parallelism with other parallelization strategies.

Prerequisites

# Install uv from https://docs.astral.sh/uv/getting-started/installation/
# Initialize the virtual environment using uv
uv venv
# Install the latest stable release from PyPI
uv pip install nemo-automodel
# Or install from source for the latest features
uv pip install git+https://github.com/NVIDIA-NeMo/Automodel.git

Before proceeding with this guide, ensure that you have NeMo AutoModel installed on your machine. For a complete guide and additional options, consult the NeMo AutoModel Installation Guide.

Key Features

AutoPipeline provides the following capabilities:

Hugging Face-Compatible Model Support: Works with decoder-only causal language models and explicitly supported VLMs that satisfy the pipeline model contract.
PyTorch Native Integration: Built on PyTorch’s torch.distributed.pipelining for optimal performance.
Flexible Configuration: Multiple scheduling strategies, configurable microbatch sizes, and automatic or manual layer splitting.
Mixed Parallelism Support: Combine pipeline parallelism with data parallelism, tensor parallelism, and FSDP.
Modular Functional API: Lower-level stage and schedule helpers can be reused with manually constructed PipelineStage objects.
Minimal Opinions: Easy to extend and integrate with existing training workflows.

Quickstart with AutoPipeline (Hugging Face Models)

Here’s a minimal example to get started with AutoPipeline using 2 pipeline stages with a Hugging Face model:

import torch
from torch.distributed.device_mesh import init_device_mesh
from nemo_automodel.components.distributed.pipelining import AutoPipeline
from nemo_automodel.components.utils.model_utils import init_empty_weights
from transformers import AutoModelForCausalLM
from transformers.initialization import no_init_weights
from transformers.utils import ContextManagers
def loss_fn(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
    """Define loss function for pipeline training."""
    return torch.nn.functional.cross_entropy(
        logits.float().view(-1, logits.size(-1)),
        targets.view(-1),
        ignore_index=-100
    )
if __name__ == "__main__":
    # 1) Initialize device mesh with 2 pipeline stages
    world_mesh = init_device_mesh("cuda", mesh_shape=(2,), mesh_dim_names=("pp",))
    # 2) Load model on meta device to avoid OOM with large models
    init_ctx = ContextManagers([no_init_weights(), init_empty_weights()])
    with init_ctx:
        model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")
    # 3) Configure and build pipeline
    ap = AutoPipeline(
        world_mesh=world_mesh,
        pp_axis_name="pp",
        pp_schedule="1f1b",
        pp_microbatch_size=1,
        pp_batch_size=8,  # Local batch size consumed by the pipeline schedule
        device=torch.cuda.current_device(),
        dtype=torch.bfloat16,
    ).build(model, loss_fn=loss_fn)
    # 4) Access pipeline components
    print(ap.debug_summary())
    print(ap.pretty_print_stages())

Run the Quickstart Example

Save the above code as pipeline_example.py and run with:

# Run with 2 GPUs for 2 pipeline stages
uv run torchrun --nproc-per-node=2 pipeline_example.py

For a complete training example:

# Run fine-tuning with 2-way pipeline parallelism using Llama 3.1 8B
automodel --nproc-per-node=2 examples/llm_finetune/llama3_1/llama3_1_8b_hellaswag_pp.yaml

Configuration Options

Basic Configuration

AutoPipeline provides comprehensive control over pipeline behavior:

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

Receive Buffer Pooling (`pp_recv_buffer_pool`)

torch.distributed.pipelining pre-allocates one receive buffer per microbatch and per direction (activations forward, gradients backward). With many microbatches and wide stage inputs those buffers can exceed the activations themselves — Kimi-K3 at 2k-token rows with 64 microbatches needs ~52 GiB of them on middle stages.

Under the 1F1B schedule only num_stages - stage_index microbatches are ever in flight on a stage, so pp_recv_buffer_pool: true rebinds the per-microbatch receive slots onto a ring of (num_stages - stage_index) + pp_recv_buffer_pool_slack real buffer sets (default slack 2). The values flowing through the schedule are unchanged; only which buffer object receives each chunk differs.

distributed:
  pipeline:
    pp_schedule: 1f1b
    pp_recv_buffer_pool: true        # default false
    pp_recv_buffer_pool_slack: 2     # extra buffer sets beyond the in-flight depth

Native interleaved1f1b also supports pp_recv_buffer_pool: true. Its pool is installed after schedule construction and uses the actual receive/compute order to assign slots separately for each local stage and direction. Received activations are copied into independent leaves before forward, and received gradients are copied before backward, so slot reuse cannot overwrite values retained by activation-checkpoint recomputation or asynchronous sends. These copies remain live for their consumers; peak memory savings therefore depend on the schedule and model rather than the ordinary 1F1B ring formula above.

Other schedules and custom CSV schedules keep the stock allocation with a warning. Ordinary 1F1B supports both known PipelineStage layouts and keeps stock behavior if neither is available. Interleaved pooling requires the native lowered communication schedule, full backward actions, and the _setup_forward_recv_info / _setup_backward_recv_info stage API. Older stage APIs keep stock buffers with a warning; unsupported lifetimes are rejected before pooling is installed.

Model Patching Flags

AutoPipeline splits a model by deep-copying it per stage and pruning away modules that do not belong to that stage. Many Hugging Face models assume the full module tree is present and return ModelOutput objects; after pruning, their original forward() often breaks (or returns objects that are awkward to pipeline).

These two flags switch AutoPipeline to lightweight, pipeline-friendly forward() implementations that return tensors (see nemo_automodel.components.distributed.pipelining.hf_utils.patch_hf_model_for_pp):

patch_inner_model: patches the decoder module (model.model for ...ForCausalLM, otherwise the module itself) so each stage can run even after pruning.
- Stage 0 (has embed_tokens): takes token IDs and produces hidden states.
- Middle stages (no embed_tokens): take hidden states from the previous stage (using inputs_embeds, or a float tensor passed through input_ids) and produce hidden states.
- Handles sliced layer containers (for example, layers becoming dict-like after stage pruning) and returns a tensor of hidden states so stages can be chained.
For compilation and performance, this patched forward prefers a precomputed causal_mask_mapping dict (it will fall back to computing masks and warn if you do not provide it).
patch_causal_lm_model: patches the ...ForCausalLM wrapper forward (the module that owns lm_head) so pipeline stages return tensors:
- Returns hidden states when lm_head is absent on that stage.
- Returns logits when lm_head is present (typically only the last stage).
- Supports logits_to_keep to compute logits for only the last k tokens.
Note: this is only used when the module you pipeline is a ...ForCausalLM-style wrapper (i.e., it has a .model attribute). If you pass a base decoder module directly, patch_causal_lm_model typically has no effect.

When to Change These Flags

Leave both True (default) for standard Hugging Face AutoModelForCausalLM / ...ForCausalLM models. This is the common case and gives the expected behavior: token IDs to hidden states to logits across stages.
Set both False when your model already has a pipeline-friendly forward. That forward returns tensors and can accept hidden states when embeddings are absent. Set both to False when the model needs custom kwargs or paths that the Hugging Face patch does not preserve. That case is common for NeMo AutoModel-native implementations, packed-sequence/thd paths, and extra arguments such as padding_mask. Many benchmark configs for NeMo-native models do this (for example, examples/llm_benchmark/qwen/qwen3_moe_30b_torch.yaml).
Set patch_inner_model=False, patch_causal_lm_model=True when your inner model is already stage-friendly, but the wrapper forward still returns a ModelOutput and you only want the wrapper simplified to “hidden states or logits”.

If you disable patch_causal_lm_model, your last stage will typically output hidden states instead of logits; in that case, make sure your loss_fn (or your last-stage module) applies the LM head explicitly.

Automatic Compared to Manual Layer Distribution

AutoPipeline offers flexible control over how your model is split across pipeline stages:

Automatic Distribution

Let AutoPipeline automatically balance layers across stages:

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

Manual Distribution

Specify exactly which modules go to each stage:

from nemo_automodel.components.distributed.pipelining.functional import (
    generate_hf_model_fqn_per_model_part
)
# Generate balanced assignments
module_fqns = generate_hf_model_fqn_per_model_part(
    num_stages=4,
    num_layers=32,
    include_embeddings=True,
    include_lm_head=True,
    include_rotary_emb=True,
    fqn_prefix="model."
)
# Or define custom assignments
custom_module_fqns = [
    # Stage 0: Embeddings + first 8 layers
    ["model.embed_tokens", "model.rotary_emb"] +
    [f"model.layers.{i}" for i in range(8)],
    # Stage 1: Next 8 layers
    ["model.rotary_emb"] + [f"model.layers.{i}" for i in range(8, 16)],
    # Stage 2: Next 8 layers
    ["model.rotary_emb"] + [f"model.layers.{i}" for i in range(16, 24)],
    # Stage 3: Final 8 layers + output
    ["model.rotary_emb"] + [f"model.layers.{i}" for i in range(24, 32)] +
    ["model.norm", "lm_head"]
]
ap = AutoPipeline(
    world_mesh=world_mesh,
    module_fqns_per_model_part=custom_module_fqns,
).build(model, loss_fn=loss_fn)

Understand Model Splitting

When AutoPipeline splits your model, it distributes components across pipeline stages. A typical split looks like this:

Example: 32-Layer Model Across 2 Stages

# Stage 0 (Rank 0): Input processing + first half
stage_0_modules = [
    "model.embed_tokens",     # Token embeddings
    "model.layers.0-15",      # First 16 transformer layers
    "model.rotary_emb"        # Position embeddings (shared)
]
# Stage 1 (Rank 1): Second half + output processing
stage_1_modules = [
    "model.layers.16-31",     # Last 16 transformer layers
    "model.norm",             # Final layer norm
    "lm_head",               # Language modeling head
    "model.rotary_emb"        # Position embeddings (shared)
]

Example: 32-Layer Model Across 4 Stages

# Stage 0 (Rank 0): Input processing
stage_0_modules = [
    "model.embed_tokens",     # Token embeddings
    "model.layers.0-7",       # First 8 transformer layers
    "model.rotary_emb"        # Position embeddings (shared)
]
# Stage 1 (Rank 1): Early layers
stage_1_modules = [
    "model.layers.8-15",      # Next 8 transformer layers
    "model.rotary_emb"
]
# Stage 2 (Rank 2): Middle layers
stage_2_modules = [
    "model.layers.16-23",     # Next 8 transformer layers
    "model.rotary_emb"
]
# Stage 3 (Rank 3): Output processing
stage_3_modules = [
    "model.layers.24-31",     # Final 8 transformer layers
    "model.norm",             # Final layer norm
    "lm_head",               # Language modeling head
    "model.rotary_emb"
]

Key observations:

Embeddings only exist on the first stage.
Language modeling head only exists on the last stage.
Rotary embeddings are shared across all stages (for position encoding).
Transformer layers are evenly distributed.

Use Lower-Level Functional Helpers

AutoPipeline and the functional pipeline_model() convenience path both use Hugging Face-specific model splitting. For a custom architecture, construct the stage modules and PyTorch PipelineStage objects in model-owned code, then reuse the lower-level scheduling helpers. The examples below illustrate those lower-level pieces; they do not make pipeline_model() a generic model splitter.

Key Functional API Components

The functional API provides several utilities for building custom pipeline parallel systems:

Stage ID Calculation

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

Module Name Generation

from nemo_automodel.components.distributed.pipelining.functional import (
    generate_hf_model_fqn_per_model_part
)
# Generate balanced module assignments for an HF-style decoder layout
module_names = generate_hf_model_fqn_per_model_part(
    num_stages=4,
    num_layers=32,
    include_embeddings=True,
    include_lm_head=True,
    include_rotary_emb=False,  # Set based on your model
    fqn_prefix=""  # Use "model." for nested models
)

Virtual Stage Calculation

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

Pipeline Schedule Build

from nemo_automodel.components.distributed.pipelining.functional import build_pipeline_schedule
# Build a schedule for your stages
schedule = build_pipeline_schedule(
    pipeline_parallel_schedule_csv=None,  # Optional CSV schedule
    pipeline_parallel_schedule="1f1b",
    microbatch_size=1,
    local_batch_size=8,
    stages=stages,  # List of PipelineStage objects
    loss_fn=loss_fn,
    scale_grads=False
)

Pseudocode: Pipeline Parallelism for Custom Models

The following fence is intentionally non-runnable pseudocode. A real custom model must provide model-owned create_stage_model splitting, a dataloader, pipeline-compatible stage input/output contracts, and process-group setup before reusing these scheduling helpers. It is an architectural outline, not a script to run unchanged.

# NON-RUNNABLE PSEUDOCODE
import torch
import torch.nn as nn
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.pipelining import PipelineStage
from nemo_automodel.components.distributed.pipelining.functional import (
    stage_ids_this_rank,
    build_pipeline_schedule,
    calculate_virtual_stages
)
class CustomTransformerBlock(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.attention = nn.MultiheadAttention(hidden_size, num_heads=8)
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.GELU(),
            nn.Linear(hidden_size * 4, hidden_size)
        )
        self.norm1 = nn.LayerNorm(hidden_size)
        self.norm2 = nn.LayerNorm(hidden_size)
    def forward(self, x):
        # Simplified transformer block
        attn_out, _ = self.attention(x, x, x)
        x = self.norm1(x + attn_out)
        x = self.norm2(x + self.mlp(x))
        return x
class CustomModel(nn.Module):
    def __init__(self, vocab_size, hidden_size, num_layers):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        self.layers = nn.ModuleList([
            CustomTransformerBlock(hidden_size) for _ in range(num_layers)
        ])
        self.output_proj = nn.Linear(hidden_size, vocab_size)
    def forward(self, input_ids):
        x = self.embedding(input_ids)
        for layer in self.layers:
            x = layer(x)
        return self.output_proj(x)
def split_custom_model_for_pipeline(model, pp_rank, pp_size, num_stages):
    """Split a custom model into pipeline stages."""
    # Determine which stages this rank handles
    stage_indices = stage_ids_this_rank(pp_rank, pp_size, num_stages, style="loop")
    stages = []
    for stage_idx in stage_indices:
        # Create a stage-specific version of the model
        # This is a simplified example - you'd need to implement proper splitting
        stage_model = create_stage_model(model, stage_idx, num_stages)
        # Create PipelineStage
        stage = PipelineStage(
            stage_model,
            stage_idx,
            num_stages,
            device=torch.cuda.current_device(),
            group=None  # Set your process group here
        )
        stages.append(stage)
    return stages
# Sketch of model-owned orchestration
def main():
    # Initialize device mesh
    world_mesh = init_device_mesh("cuda", mesh_shape=(4,), mesh_dim_names=("pp",))
    pp_rank = world_mesh["pp"].get_local_rank()
    pp_size = world_mesh["pp"].size()
    # Create model
    model = CustomModel(vocab_size=50000, hidden_size=768, num_layers=24)
    # Calculate virtual stages
    num_virtual_stages, stages_per_rank = calculate_virtual_stages(
        num_layers=24,
        layers_per_stage=3,  # 8 virtual stages total
        pp_size=4,
        is_single_stage_schedule=False
    )
    # Split model into stages
    stages = split_custom_model_for_pipeline(model, pp_rank, pp_size, num_virtual_stages)
    # Define loss function
    def loss_fn(logits, targets):
        return nn.functional.cross_entropy(
            logits.view(-1, logits.size(-1)),
            targets.view(-1)
        )
    # Build pipeline schedule
    schedule = build_pipeline_schedule(
        pipeline_parallel_schedule_csv=None,
        pipeline_parallel_schedule="interleaved1f1b",  # Multi-stage schedule registered by PyTorch
        microbatch_size=1,
        local_batch_size=8,
        stages=stages,
        loss_fn=loss_fn,
        scale_grads=True
    )
    # Training loop
    for batch in dataloader:
        # Use schedule.step() for training
        losses = []
        schedule.step(batch["input_ids"], target=batch["labels"], losses=losses)
        # losses will contain the loss values from the last stage
        if losses:
            print(f"Loss: {sum(losses) / len(losses)}")

Add a Parallelization Callback to Hugging Face-Compatible Splitting

When using the Hugging Face-compatible pipeline_model() path, a callback can apply additional parallelism to each model part. Its signature must match ParallelizeFnProtocol:

from nemo_automodel.components.distributed.pipelining.functional import pipeline_model
def custom_parallelize_fn(
    model, world_mesh, moe_mesh, *,
    dp_axis_names,
    cp_axis_name=None,
    tp_axis_name=None,
    ep_axis_name=None,
    ep_shard_axis_names=None,
):
    """Custom parallelization function for each local model part."""
    # Apply your custom parallelization logic here
    # This is called for each local model part
    if dp_axis_names:
        # Apply data parallelism
        pass
    # Add any other parallelization strategies
    pass
# Use pipeline_model only with an HF-compatible model.
schedule, model_parts, has_first, has_last, stages = pipeline_model(
    model=hf_compatible_model,
    world_mesh=world_mesh,
    moe_mesh=None,
    pp_axis_name="pp",
    dp_axis_names=("dp",),
    layers_per_stage=4,
    pipeline_parallel_schedule="1f1b",
    pipeline_parallel_schedule_csv=None,
    microbatch_size=1,
    local_batch_size=8,
    device=torch.cuda.current_device(),
    loss_fn=loss_fn,
    parallelize_fn=custom_parallelize_fn,
    module_fqns_per_model_part=None,  # Or provide an explicit HF module assignment
    patch_inner_model=False,  # Use only when the model's inner forward is already stage-friendly
    patch_causal_lm_model=False,  # Use only when the wrapper already returns stage tensors
)

Tips for Manually Staged Custom Models

The lower-level helpers can support a custom model after its model-owned splitting code has created the stage modules:

Module Naming: Ensure your model has consistent module naming that can be mapped to stages.
State Management: Handle model state (embeddings, buffers) carefully across stages.
Communication: First and last stages need special handling for inputs and outputs.
Splitting Ownership: Keep custom model splitting in the model implementation; do not pass an arbitrary nn.Module to pipeline_model().
Testing: Start with a small model and verify correct splitting before scaling up.

The schedule helpers are reusable, but the repository’s automatic splitting path retains Hugging Face-specific assumptions.

Mixed Parallelism

AutoPipeline can be combined with other parallelization strategies for optimal performance:

def parallelize_fn(
    model, world_mesh, moe_mesh, *,
    dp_axis_names,
    cp_axis_name=None,
    tp_axis_name=None,
    ep_axis_name=None,
    ep_shard_axis_names=None,
):
    """Apply additional parallelization to each local model part."""
    # Example: Apply FSDP to each local model part
    if dp_axis_names:
        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
        # Wrap model with FSDP (simplified example)
        # In practice, you'd configure FSDP parameters
        pass
    # Example: Apply tensor parallelism
    if tp_axis_name:
        # Apply tensor parallelism to attention/MLP layers
        pass
# Build pipeline with custom parallelization
ap = AutoPipeline(world_mesh=world_mesh).build(
    model,
    loss_fn=loss_fn,
    parallelize_fn=parallelize_fn
)

Monitor and Debug

AutoPipeline provides comprehensive tools for understanding your pipeline configuration:

Pipeline Information

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

Analysis

# Parameter distribution
stage_param_counts = ap.get_stage_param_counts()
total_params = ap.get_total_param_count()
trainable_params = ap.get_total_param_count(trainable_only=True)
for i, params in enumerate(stage_param_counts):
    percentage = (params / total_params) * 100
    print(f"Stage {i}: {params:,} parameters ({percentage:.1f}%)")
# Debug summary
print(ap.debug_summary())
print(ap.pretty_print_stages(max_modules_per_stage=10))
# Visualize schedule
ap.visualize_current_schedule("pipeline_schedule.png")

Gradient Management

from nemo_automodel.components.distributed.pipelining.functional import scale_grads_by_divisor
from nemo_automodel.components.training.utils import clip_grad_norm
# Call these after backward on every pipeline rank.
scale_grads_by_divisor(ap.info.stages, divisor=8)
# Use the PP-aware clipping helper across the local model parts. For MoE
# models with expert parallelism, also pass moe_mesh/ep_axis_name so expert
# gradient norms are reduced across the EP axis.
grad_norm = clip_grad_norm(
    1.0,
    ap.parts,
    norm_type=2.0,
    pp_enabled=True,
    device_mesh=world_mesh,
    pp_axis_name="pp",
    moe_mesh=moe_mesh,
    ep_axis_name="ep" if moe_mesh is not None and "ep" in (moe_mesh.mesh_dim_names or ()) else None,
)

Add Pipeline Parallelism to Supported Configurations

Pipeline parallelism can be added to an FSDP2 training configuration only when the model satisfies the repository’s PP contract and the PP, TP, CP, and inferred DP sizes compose to the total worker count. DDP and Megatron FSDP do not support PP. For recipe-driven model implementations, use a model with a _pp_plan or the explicitly supported MoE path.

Command-Line Override Method

For an FSDP2 config whose model supports PP, use command-line overrides such as:

uv run automodel examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml \
    --nproc-per-node 2 \
    --step_scheduler.local_batch_size 8 \
    --distributed.strategy fsdp2 \
    --distributed.pp_size 2 \
    --distributed.pipeline.pp_schedule 1f1b \
    --distributed.pipeline.pp_microbatch_size 1 \
    --distributed.pipeline.round_virtual_stages_to_pp_multiple up \
    --distributed.pipeline.scale_grads_in_schedule false

Key parameters to override:

--distributed.pp_size: Number of pipeline ranks; PP, TP, CP, and inferred DP sizes must compose to the total worker count
--step_scheduler.local_batch_size: Runtime batch size used for both the pipeline schedule and dataloader
--distributed.pipeline.pp_schedule: Pipeline schedule (1f1b, interleaved1f1b, LoopedBFS, and so on)

YAML Configuration Method

Add these sections to a compatible FSDP2 YAML config:

step_scheduler:
  local_batch_size: 8  # Also sets the runtime pipeline batch size
distributed:
  strategy: fsdp2
  dp_size: 1
  tp_size: 1
  cp_size: 1
  pp_size: 4  # Enable 4-way pipeline parallelism
  sequence_parallel: false
  pipeline:
    pp_schedule: 1f1b
    pp_microbatch_size: 1
    round_virtual_stages_to_pp_multiple: up
    scale_grads_in_schedule: false
    layers_per_stage: null  # Auto-compute, or specify number

Mixed Parallelism Examples

Pipeline and Data Parallelism (4 GPUs Total)

uv run automodel your_config.yaml \
    --nproc-per-node 4 \
    --distributed.pp_size 2 \
    --distributed.dp_size 2 \
    --step_scheduler.local_batch_size 16

Pipeline and Tensor Parallelism (4 GPUs Total)

uv run automodel your_config.yaml \
    --nproc-per-node 4 \
    --distributed.pp_size 2 \
    --distributed.tp_size 2 \
    --step_scheduler.local_batch_size 8

Full Hybrid with PP, DP, and TP (8 GPUs Total)

uv run automodel your_config.yaml \
    --nproc-per-node 8 \
    --distributed.pp_size 2 \
    --distributed.dp_size 2 \
    --distributed.tp_size 2 \
    --step_scheduler.local_batch_size 32

Integrate with Training Recipes

AutoPipeline integrates with NeMo AutoModel’s recipe system. The following YAML configuration is a complete example:

# config.yaml
recipe: TrainFinetuneRecipeForNextTokenPrediction
step_scheduler:
  global_batch_size: 8
  local_batch_size: 8
  max_steps: 10
  ckpt_every_steps: 10
  val_every_steps: 10
dist_env:
  backend: nccl
  timeout_minutes: 1
seed: 1111
model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
checkpoint:
  enabled: false
clip_grad_norm:
  max_norm: 1.0
distributed:
  strategy: fsdp2
  dp_size: none
  dp_replicate_size: none
  tp_size: 1
  cp_size: 1
  pp_size: 2          # 2-way pipeline parallelism
  ep_size: 1
  sequence_parallel: false
  pipeline:
    pp_schedule: 1f1b
    pp_microbatch_size: 1
    layers_per_stage: null  # Auto-compute layer distribution
    round_virtual_stages_to_pp_multiple: up
    scale_grads_in_schedule: false
loss_fn:
  _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy
dataset:
  _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
  dataset_name: rajpurkar/squad
  split: train
packed_sequence:
  packed_sequence_size: 0
dataloader:
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
  collate_fn:
    _target_: nemo_automodel.components.datasets.utils.default_collater
  shuffle: false
validation_dataset:
  _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
  dataset_name: rajpurkar/squad
  split: validation
  limit_dataset_samples: 64
validation_dataloader:
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
  collate_fn:
    _target_: nemo_automodel.components.datasets.utils.default_collater
optimizer:
  _target_: torch.optim.Adam
  betas: [0.9, 0.999]
  eps: 1e-8
  lr: 1.0e-5
  weight_decay: 0
lr_scheduler:
  lr_decay_style: cosine
  min_lr: 1.0e-6

Run training with:

# Run with 2 GPUs for 2-way pipeline parallelism
uv run automodel config.yaml --nproc-per-node 2

Troubleshooting

Common Issues

The model does not fit in memory:

Increase the number of pipeline stages.
Reduce microbatch size.
Enable gradient checkpointing.

Pipeline bubbles reducing efficiency:

Increase step_scheduler.local_batch_size to have more microbatches.
Try different schedules (for example, interleaved1f1b).
Adjust virtual stages configuration.

Uneven stage distribution:

Use manual module assignment for fine control.
Adjust layers_per_stage parameter.
Check parameter counts with get_stage_param_counts().

Conclusion

AutoPipeline and pipeline_model() provide the repository’s Hugging Face-compatible splitting path, while lower-level functional helpers can be reused after custom code constructs its own pipeline stages.

Key takeaways:

Pipeline parallelism enables training of models too large for a single GPU.
AutoPipeline provides a simple API for Hugging Face-compatible models with customization options.
Lower-level functional helpers can schedule manually constructed stages for custom models.
Both can be combined with other parallelization strategies for optimal performance.
Use built-in monitoring tools to understand and optimize your pipeline.