nemo_automodel.components.distributed.pipelining.runtime

Module Contents

Classes

Name	Description
`PipelineRuntimeInitializer`	A model-owned runtime resource that must be ready before a PP step.

Functions

Name	Description
`collect_pipeline_runtime_initializers`	Collect model-owned initializers, deduplicating compatible shared resources.

API

class nemo_automodel.components.distributed.pipelining.runtime.PipelineRuntimeInitializer()

Protocol

A model-owned runtime resource that must be ready before a PP step.

resource_key

Hashable

Identity of the process-local runtime resource this initializer owns.

signature

Hashable

Immutable configuration expected by resource_key.

nemo_automodel.components.distributed.pipelining.runtime.PipelineRuntimeInitializer.prepare(
    num_tokens: int,
    device: torch.device
) -> None

Prepare the resource for the upcoming pipeline microbatch.

Parameters:

num_tokens

int

Maximum number of tokens in the upcoming pipeline microbatch.

device

torch.device

Device used by the pipeline stage.

nemo_automodel.components.distributed.pipelining.runtime.collect_pipeline_runtime_initializers(
    model_parts: list[torch.nn.Module]
) -> list[nemo_automodel.components.distributed.pipelining.runtime.PipelineRuntimeInitializer]

Collect model-owned initializers, deduplicating compatible shared resources.

The pipeline layer treats keys and signatures as opaque values. A repeated resource key is initialized once when every provider reports the same signature; conflicting signatures are rejected before any pipeline traffic.

Parameters:

model_parts

list[nn.Module]

Pipeline-local model partitions whose modules may provide runtime initializers.

Returns: list[PipelineRuntimeInitializer]

Initializers in module traversal order, with compatible shared resources deduplicated.

Raises:

TypeError: If an initializer’s resource key or signature is not hashable.
RuntimeError: If providers report incompatible signatures for the same resource key.