nemo_automodel.components.distributed.pipelining.config

Pipeline parallel configuration class.

Design principle:

Device mesh (world_mesh, moe_mesh) is passed separately to from_pretrained/from_config
PipelineConfig contains scheduling, splitting, and runtime options
loss_fn is included here since it’s only used for pipelining
Axis names are inferred automatically from device_mesh in _instantiate_pipeline

Usage: from nemo_automodel.components.distributed.pipelining.config import PipelineConfig

config = PipelineConfig( pp_schedule=“1f1b”, pp_microbatch_size=2, pp_batch_size=8, loss_fn=my_loss_fn, )

Module Contents

Classes

Name	Description
`PipelineConfig`	Configuration for pipeline parallel training.

API

class nemo_automodel.components.distributed.pipelining.config.PipelineConfig(
    pp_schedule: str | None = '1f1b',
    pp_schedule_csv: str | None = None,
    pp_microbatch_size: int = 1,
    pp_batch_size: int = 1,
    layers_per_stage: int | None = None,
    round_virtual_stages_to_pp_multiple: typing.Literal['up', 'down'] | None = None,
    module_fqns_per_model_part: typing.List[typing.List[str]] | None = None,
    patch_inner_model: bool = True,
    patch_causal_lm_model: bool = True,
    patch_stage_backward_maybe_with_nosync: bool = False,
    dtype: torch.dtype | None = None,
    scale_grads_in_schedule: bool = False,
    loss_fn: typing.Callable | None = None,
    pp_seq_len: int | None = None,
    pp_recv_buffer_pool: bool = False,
    pp_recv_buffer_pool_slack: int = 2
)

Dataclass

Configuration for pipeline parallel training.

Note: Device mesh (world_mesh, moe_mesh) is passed separately on the from_pretrained/from_config method signature. Pipeline parallelism is enabled when pp_size > 1. Axis names are inferred automatically from the device mesh structure.

dtype

dtype | None = None

Data type for pipeline computation. If None, uses the model’s default dtype.

layers_per_stage

int | None = None

Number of transformer layers per pipeline stage. If None, layers are split evenly across stages.

loss_fn

Callable | None = None

Loss function used for pipeline training. Required when pipeline is enabled. The function signature should be compatible with the model’s output format.

module_fqns_per_model_part

List[List[str]] | None = None

Explicit specification of which module FQNs belong to each model part/stage. If provided, overrides automatic layer splitting.

patch_causal_lm_model

bool = True

Apply pipeline patches to the CausalLM wrapper model. Defaults to True.

patch_inner_model

bool = True

Apply pipeline patches to the inner model (e.g., the base transformer in a CausalLM wrapper). Defaults to True.

patch_stage_backward_maybe_with_nosync

bool = False

Patch stage backward to use no_sync context for gradient accumulation efficiency. Useful when combining PP with FSDP.

pp_batch_size

int = 1

Total batch size per pipeline stage. Must be divisible by pp_microbatch_size.

pp_microbatch_size

int = 1

Size of each microbatch for pipeline execution. pp_batch_size must be divisible by pp_microbatch_size.

pp_recv_buffer_pool

bool = False

Pool the per-microbatch P2P recv buffers into a ring sized to the 1F1B in-flight depth instead of pre-allocating one buffer set per microbatch and direction. Cuts pipeline-stage buffer memory roughly by num_microbatches / (num_stages - stage_index + slack) on early/middle stages. For native interleaved 1F1B, receives are pooled using the lowered schedule’s actual lifetimes. Inputs are copied into independent leaves before forward, and gradients are copied before backward to protect checkpoint recomputation and escaping aliases. Other schedules are ignored with a warning. Defaults to False.

pp_recv_buffer_pool_slack

int = 2

Extra buffer sets beyond the 1F1B in-flight depth when pp_recv_buffer_pool is enabled. Defaults to 2.

pp_schedule

str | None = '1f1b'

Pipeline schedule type. Supported values: “1f1b” (one-forward-one-backward), “gpipe”, “interleaved_1f1b”, “looped_bfs”, “dfs”, “v_schedule”, “zero_bubble”. Defaults to “1f1b”.

pp_schedule_csv

str | None = None

Path to a CSV file defining a custom pipeline schedule. If provided, overrides pp_schedule.

pp_seq_len

int | None = None

Sequence length hint for pipeline parallel shape inference. If None, it will be inferred from the dataset config.

round_virtual_stages_to_pp_multiple

Literal['up', 'down'] | None = None

When using virtual stages (interleaved schedules), round the number of virtual stages to a multiple of pp_size. “up” rounds up, “down” rounds down. If None, no rounding is applied.

scale_grads_in_schedule

bool = False

Scale gradients within the pipeline schedule (by 1/n_microbatches). If False, gradients must be scaled externally. Defaults to False.

nemo_automodel.components.distributed.pipelining.config.PipelineConfig.__post_init__() -> None

Validate the recv-buffer-pool knobs.

nemo_automodel.components.distributed.pipelining.config.PipelineConfig.to_dict() -> typing.Dict[str, typing.Any]

Convert config to dictionary.