nemo_automodel.components.cuda_graphs.partial

Component-owned partial CUDA graphs for fixed-shape attention and MoE preprocessing.

Attention, the parameterless router core, and HybridEP metadata preprocessing may be captured. Dynamic dispatch, expert compute, and combine remain eager.

Module Contents

Classes

Name	Description
`PartialCudaGraphManager`	Capture selected TE attention and MoE submodules after one eager iteration.
`_CapturedCall`	Detached sample inputs and the invariants required for a safe replay.
`_DiscoveredBlock`	One structurally discovered transformer or MTP block.
`_ExplicitParameterCallAdapter`	Present module parameters as graph inputs instead of captured module state.
`_GraphModuleSpec`	Discovery and entry-construction contract for one user-facing graph module.
`_PartialGraphEntry`	One graphable module invocation and its replay safety state.
`_TensorOnlyCallAdapter`	Expose a mixed tensor/control module call as a tensor-only module call.

Functions

Name	Description
`_alias_pattern`	Encode repeated tensor objects without depending on their absolute identities.
`_build_moe_preprocess_entry`	Build one fused HybridEP metadata-preprocessing graph entry.
`_build_moe_router_entry`	Build one parameterless MoE router graph entry.
`_build_te_dpa_entry`	Build one TE fused-attention graph entry.
`_build_whole_attention_entry`	Build one whole-attention graph entry with graph-safe explicit parameters.
`_canonicalize_bf16_fused_attention`	Remove ignored FP8 metadata while requiring the DPA compute itself to be BF16.
`_describe_control_value`	Return a bounded diagnostic for one non-tensor graph control.
`_discover_blocks`	Discover main-stack and MTP blocks using shared structural traversal.
`_find_attention`	Find the whole attention module in a transformer block.
`_find_moe_module`	Find an MoE sublayer by its gate-and-experts capability.
`_find_moe_preprocess`	Find graphable fused HybridEP metadata preprocessing.
`_find_moe_router`	Find the parameterless routing core in an MoE transformer block.
`_find_te_dpa`	Find the parameterless TE fused-attention boundary in a transformer block.
`_get_make_graphed_callables`	Load Transformer Engine’s graph helper only when the feature is enabled.
`_is_transformer_engine_pybind_enum`	Return whether `value` is an immutable enum exported by TE’s extension.
`_model_label`	Return a diagnostic model label without using it as a capability gate.
`_named_buffer_storage`	Return buffer identities and storage properties captured by a graph.
`_require_local_parameter_storage`	Reject sharded or otherwise unmaterialized parameters before graph capture.
`_require_parameterless_target`	Reject graph boundaries that own parameters managed outside the graph.
`_same_control_value`	Compare non-tensor graph controls without invoking tensor-like equality.
`_select_graph_targets`	Select graph targets while rejecting shared/repeated physical call sites.
`_tensor_metadata`	Return metadata that must stay invariant across graph replays.
`_unwrap_checkpoint_wrappers`	Return the module beneath any nested PyTorch checkpoint wrappers.
`_uses_repeated_mtp_layer`	Return whether one physical MTP layer is invoked repeatedly per forward.

Data

_Canonicalizer

_GRAPH_MODULE_SPECS

logger

API

class nemo_automodel.components.cuda_graphs.partial.PartialCudaGraphManager(
    entries: list[nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry]
)

Capture selected TE attention and MoE submodules after one eager iteration.

nemo_automodel.components.cuda_graphs.partial.PartialCudaGraphManager.capture() -> None

Batch-capture all observed targets in real forward order.

nemo_automodel.components.cuda_graphs.partial.PartialCudaGraphManager.close() -> None

Idempotently destroy every partial graph before distributed teardown.

nemo_automodel.components.cuda_graphs.partial.PartialCudaGraphManager.from_model_parts(
    model_parts: list[torch.nn.Module],
    activation_checkpointing: bool = False,
    pipeline_parallel: bool = False
) -> nemo_automodel.components.cuda_graphs.partial.PartialCudaGraphManager | None

classmethod

Discover graph targets from an already-built training model.

nemo_automodel.components.cuda_graphs.partial.PartialCudaGraphManager.log_stats(
    phase: str
) -> None

Log visible aggregate graph activity counters.

nemo_automodel.components.cuda_graphs.partial.PartialCudaGraphManager.start_recording() -> None

Install first-call recorders on every selected target.

nemo_automodel.components.cuda_graphs.partial.PartialCudaGraphManager.stats() -> dict[str, int]

Return aggregate capture, replay, and eager-fallback counters.

class nemo_automodel.components.cuda_graphs.partial._CapturedCall(
    tree_spec: typing.Any,
    template_leaves: tuple[typing.Any, ...],
    tensor_positions: tuple[int, ...],
    tensor_input_indices: tuple[int, ...],
    sample_tensors: tuple[torch.Tensor, ...],
    tensor_metadata: tuple[tuple[typing.Any, ...], ...]
)

Dataclass

Detached sample inputs and the invariants required for a safe replay.

sample_tensors

tuple[Tensor, ...]

template_leaves

tuple[Any, ...]

tensor_input_indices

tuple[int, ...]

tensor_metadata

tuple[tuple[Any, ...], ...]

tensor_positions

tuple[int, ...]

nemo_automodel.components.cuda_graphs.partial._CapturedCall.from_call(
    args: tuple[typing.Any, ...],
    kwargs: dict[str, typing.Any]
) -> nemo_automodel.components.cuda_graphs.partial._CapturedCall

classmethod

nemo_automodel.components.cuda_graphs.partial._CapturedCall.rebuild(
    tensors: collections.abc.Sequence[torch.Tensor]
) -> tuple[tuple[typing.Any, ...], dict[str, typing.Any]]

Reconstruct the target call and its aliases from unique tensor inputs.

nemo_automodel.components.cuda_graphs.partial._CapturedCall.validate(
    args: tuple[typing.Any, ...],
    kwargs: dict[str, typing.Any]
) -> tuple[bool, str, tuple[torch.Tensor, ...]]

Validate a replay call and return its flattened tensor inputs.

class nemo_automodel.components.cuda_graphs.partial._DiscoveredBlock(
    name: str,
    module: torch.nn.Module,
    capture_owner: torch.nn.Module | None,
    moe: torch.nn.Module | None,
    is_mtp: bool
)

Dataclass

One structurally discovered transformer or MTP block.

capture_owner

Module | None

is_mtp

bool

module

Module

moe

Module | None

name

str

class nemo_automodel.components.cuda_graphs.partial._ExplicitParameterCallAdapter(
    target: torch.nn.Module,
    captured_call: nemo_automodel.components.cuda_graphs.partial._CapturedCall
)

Bases: Module

Present module parameters as graph inputs instead of captured module state.

buffer_storage

= _named_buffer_storage(target)

capture_inputs

tuple[Tensor, ...]

Return dynamic samples followed by graph-owned parameter samples.

capture_parameters

= tuple(capture_parameters)

dynamic_input_count

= len(captured_call.sample_tensors)

parameter_metadata

parameter_names

nemo_automodel.components.cuda_graphs.partial._ExplicitParameterCallAdapter.forward(
    tensor_inputs: torch.Tensor = ()
) -> typing.Any

Run the target with explicit parameter values.

nemo_automodel.components.cuda_graphs.partial._ExplicitParameterCallAdapter.named_parameters(
    prefix: str = '',
    recurse: bool = True,
    remove_duplicate: bool = True
) -> collections.abc.Iterator[tuple[str, torch.nn.Parameter]]

Hide the registered target parameters from generic module utilities.

nemo_automodel.components.cuda_graphs.partial._ExplicitParameterCallAdapter.parameters(
    recurse: bool = True
) -> collections.abc.Iterator[torch.nn.Parameter]

Hide the registered target parameters from TE module-parameter discovery.

nemo_automodel.components.cuda_graphs.partial._ExplicitParameterCallAdapter.replay_inputs(
    dynamic_inputs: tuple[torch.Tensor, ...]
) -> tuple[torch.Tensor, ...]

Append the currently materialized live parameters for graph replay.

class nemo_automodel.components.cuda_graphs.partial._GraphModuleSpec(
    find_target: collections.abc.Callable[[_DiscoveredBlock], torch.nn.Module | None],
    build_entry: collections.abc.Callable[[_DiscoveredBlock, nn.Module, bool], nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry]
)

Dataclass

Discovery and entry-construction contract for one user-facing graph module.

build_entry

Callable[[_DiscoveredBlock, nn.Module, bool], _PartialGraphEntry]

find_target

Callable[[_DiscoveredBlock], Module | None]

class nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry(
    name: str,
    target: torch.nn.Module,
    canonicalizer: nemo_automodel.components.cuda_graphs.partial._Canonicalizer | None = None,
    capture_input_variants: int = 1,
    explicit_parameters: bool = False,
    capture_owner: torch.nn.Module | None = None,
    retain_graph_in_backward: bool = False
)

One graphable module invocation and its replay safety state.

_adapters

tuple[Module, ...] = ()

_captured_call_variants

list[_CapturedCall] = []

_captured_training

bool | None = None

capture_count

= 0

captured_call

_CapturedCall | None = None

fallback_count

= 0

original_forward

= target.forward

replay_count

= 0

nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry._canonicalize(
    args: tuple[typing.Any, ...],
    kwargs: dict[str, typing.Any]
) -> tuple[tuple[typing.Any, ...], dict[str, typing.Any]]

nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry.build_adapter(
    captured_call: nemo_automodel.components.cuda_graphs.partial._CapturedCall | None = None
) -> torch.nn.Module

Build the tensor-only capture adapter after an eager sample was observed.

nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry.capture_inputs(
    adapter: torch.nn.Module,
    captured_call: nemo_automodel.components.cuda_graphs.partial._CapturedCall
) -> tuple[torch.Tensor, ...]

Return the sample surface passed to Transformer Engine.

nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry.captured_calls() -> tuple[nemo_automodel.components.cuda_graphs.partial._CapturedCall, ...]

Return the distinct iteration-0 input contracts in call order.

nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry.close() -> None

Destroy graph replay state and restore the eager target.

nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry.finish_capture() -> None

Return an FSDP2 owner to its normal sharded state after capture.

nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry.install(
    graphed_adapter: torch.nn.Module | collections.abc.Sequence[torch.nn.Module]
) -> None

Install validated graph replay around the original target forward.

nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry.prepare_for_capture() -> None

Materialize FSDP2 parameters before constructing explicit graph inputs.

nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry.start_recording() -> None

Record the first eager call through a temporary pre-hook.

nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry.stop_recording() -> None

Remove the temporary eager-call recorder.

class nemo_automodel.components.cuda_graphs.partial._TensorOnlyCallAdapter(
    target: torch.nn.Module,
    captured_call: nemo_automodel.components.cuda_graphs.partial._CapturedCall
)

Bases: Module

Expose a mixed tensor/control module call as a tensor-only module call.

nemo_automodel.components.cuda_graphs.partial._TensorOnlyCallAdapter.forward(
    tensor_inputs: torch.Tensor = ()
) -> typing.Any

Rebuild and execute the captured target call.

nemo_automodel.components.cuda_graphs.partial._alias_pattern(
    tensors: collections.abc.Sequence[torch.Tensor]
) -> tuple[int, ...]

Encode repeated tensor objects without depending on their absolute identities.

nemo_automodel.components.cuda_graphs.partial._build_moe_preprocess_entry(
    block: nemo_automodel.components.cuda_graphs.partial._DiscoveredBlock,
    target: torch.nn.Module,
    _activation_checkpointing: bool
) -> nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry

Build one fused HybridEP metadata-preprocessing graph entry.

nemo_automodel.components.cuda_graphs.partial._build_moe_router_entry(
    block: nemo_automodel.components.cuda_graphs.partial._DiscoveredBlock,
    target: torch.nn.Module,
    _activation_checkpointing: bool
) -> nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry

Build one parameterless MoE router graph entry.

nemo_automodel.components.cuda_graphs.partial._build_te_dpa_entry(
    block: nemo_automodel.components.cuda_graphs.partial._DiscoveredBlock,
    target: torch.nn.Module,
    activation_checkpointing: bool
) -> nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry

Build one TE fused-attention graph entry.

nemo_automodel.components.cuda_graphs.partial._build_whole_attention_entry(
    block: nemo_automodel.components.cuda_graphs.partial._DiscoveredBlock,
    target: torch.nn.Module,
    _activation_checkpointing: bool
) -> nemo_automodel.components.cuda_graphs.partial._PartialGraphEntry

Build one whole-attention graph entry with graph-safe explicit parameters.

nemo_automodel.components.cuda_graphs.partial._canonicalize_bf16_fused_attention(
    args: tuple[typing.Any, ...],
    kwargs: dict[str, typing.Any]
) -> tuple[tuple[typing.Any, ...], dict[str, typing.Any]]

Remove ignored FP8 metadata while requiring the DPA compute itself to be BF16.

nemo_automodel.components.cuda_graphs.partial._describe_control_value(
    value: typing.Any
) -> str

Return a bounded diagnostic for one non-tensor graph control.

nemo_automodel.components.cuda_graphs.partial._discover_blocks(
    model: torch.nn.Module
) -> list[nemo_automodel.components.cuda_graphs.partial._DiscoveredBlock]

Discover main-stack and MTP blocks using shared structural traversal.

nemo_automodel.components.cuda_graphs.partial._find_attention(
    block: nemo_automodel.components.cuda_graphs.partial._DiscoveredBlock
) -> torch.nn.Module | None

Find the whole attention module in a transformer block.

nemo_automodel.components.cuda_graphs.partial._find_moe_module(
    block: torch.nn.Module
) -> torch.nn.Module | None

Find an MoE sublayer by its gate-and-experts capability.

nemo_automodel.components.cuda_graphs.partial._find_moe_preprocess(
    block: nemo_automodel.components.cuda_graphs.partial._DiscoveredBlock
) -> torch.nn.Module | None

Find graphable fused HybridEP metadata preprocessing.

nemo_automodel.components.cuda_graphs.partial._find_moe_router(
    block: nemo_automodel.components.cuda_graphs.partial._DiscoveredBlock
) -> torch.nn.Module | None

Find the parameterless routing core in an MoE transformer block.

nemo_automodel.components.cuda_graphs.partial._find_te_dpa(
    block: nemo_automodel.components.cuda_graphs.partial._DiscoveredBlock
) -> torch.nn.Module | None

Find the parameterless TE fused-attention boundary in a transformer block.

nemo_automodel.components.cuda_graphs.partial._get_make_graphed_callables() -> collections.abc.Callable[..., typing.Any]

Load Transformer Engine’s graph helper only when the feature is enabled.

nemo_automodel.components.cuda_graphs.partial._is_transformer_engine_pybind_enum(
    value: typing.Any
) -> bool

Return whether value is an immutable enum exported by TE’s extension.

nemo_automodel.components.cuda_graphs.partial._model_label(
    model: torch.nn.Module
) -> str

Return a diagnostic model label without using it as a capability gate.

nemo_automodel.components.cuda_graphs.partial._named_buffer_storage(
    target: torch.nn.Module
) -> tuple[tuple[typing.Any, ...], ...]

Return buffer identities and storage properties captured by a graph.

nemo_automodel.components.cuda_graphs.partial._require_local_parameter_storage(
    name: str,
    parameter: torch.nn.Parameter
) -> None

Reject sharded or otherwise unmaterialized parameters before graph capture.

nemo_automodel.components.cuda_graphs.partial._require_parameterless_target(
    module_name: str,
    block: nemo_automodel.components.cuda_graphs.partial._DiscoveredBlock,
    target: torch.nn.Module
) -> None

Reject graph boundaries that own parameters managed outside the graph.

nemo_automodel.components.cuda_graphs.partial._same_control_value(
    expected: typing.Any,
    actual: typing.Any
) -> bool

Compare non-tensor graph controls without invoking tensor-like equality.

nemo_automodel.components.cuda_graphs.partial._select_graph_targets(
    module_name: str,
    candidates: list[tuple[nemo_automodel.components.cuda_graphs.partial._DiscoveredBlock, torch.nn.Module]],
    repeated_mtp_layer: bool
) -> dict[str, torch.nn.Module]

Select graph targets while rejecting shared/repeated physical call sites.

nemo_automodel.components.cuda_graphs.partial._tensor_metadata(
    tensor: torch.Tensor
) -> tuple[typing.Any, ...]

Return metadata that must stay invariant across graph replays.

Parameters:

tensor

torch.Tensor

Tensor of arbitrary shape whose type, layout, shape, strides, dtype, device, and autograd requirement define the replay contract.

Returns: Any

Tuple containing the tensor subclass and replay-critical metadata. The

nemo_automodel.components.cuda_graphs.partial._unwrap_checkpoint_wrappers(
    module: torch.nn.Module
) -> torch.nn.Module

Return the module beneath any nested PyTorch checkpoint wrappers.

nemo_automodel.components.cuda_graphs.partial._uses_repeated_mtp_layer(
    model: torch.nn.Module
) -> bool

Return whether one physical MTP layer is invoked repeatedly per forward.

nemo_automodel.components.cuda_graphs.partial._Canonicalizer = Callable[[tuple[Any, ...], dict[str, Any]], tuple[tuple[Any, ...], dict[str, Any...

nemo_automodel.components.cuda_graphs.partial._GRAPH_MODULE_SPECS = {'attn': _GraphModuleSpec(_find_attention, _build_whole_attention_entry), 'te_dp...

nemo_automodel.components.cuda_graphs.partial.logger = logging.getLogger(__name__)