> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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`](#nemo_automodel-components-cuda_graphs-partial-PartialCudaGraphManager)             | Capture selected TE attention and MoE submodules after one eager iteration. |
| [`_CapturedCall`](#nemo_automodel-components-cuda_graphs-partial-_CapturedCall)                                 | Detached sample inputs and the invariants required for a safe replay.       |
| [`_DiscoveredBlock`](#nemo_automodel-components-cuda_graphs-partial-_DiscoveredBlock)                           | One structurally discovered transformer or MTP block.                       |
| [`_ExplicitParameterCallAdapter`](#nemo_automodel-components-cuda_graphs-partial-_ExplicitParameterCallAdapter) | Present module parameters as graph inputs instead of captured module state. |
| [`_GraphModuleSpec`](#nemo_automodel-components-cuda_graphs-partial-_GraphModuleSpec)                           | Discovery and entry-construction contract for one user-facing graph module. |
| [`_PartialGraphEntry`](#nemo_automodel-components-cuda_graphs-partial-_PartialGraphEntry)                       | One graphable module invocation and its replay safety state.                |
| [`_TensorOnlyCallAdapter`](#nemo_automodel-components-cuda_graphs-partial-_TensorOnlyCallAdapter)               | Expose a mixed tensor/control module call as a tensor-only module call.     |

### Functions

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

### Data

[`_Canonicalizer`](#nemo_automodel-components-cuda_graphs-partial-_Canonicalizer)

[`_GRAPH_MODULE_SPECS`](#nemo_automodel-components-cuda_graphs-partial-_GRAPH_MODULE_SPECS)

[`logger`](#nemo_automodel-components-cuda_graphs-partial-logger)

### API

```python
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.

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

Batch-capture all observed targets in real forward order.

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

Idempotently destroy every partial graph before distributed teardown.

```python
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.

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

Log visible aggregate graph activity counters.

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

Install first-call recorders on every selected target.

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

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

```python
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, ...]`

---

```python
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

```python
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.

```python
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.

```python
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`

---

```python
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`**

---

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

Run the target with explicit parameter values.

```python
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.

```python
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.

```python
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.

```python
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]`

---

```python
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`

---

```python
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]]
```

```python
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.

```python
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.

```python
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.

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

Destroy graph replay state and restore the eager target.

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

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

```python
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.

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

Materialize FSDP2 parameters before constructing explicit graph inputs.

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

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

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

Remove the temporary eager-call recorder.

```python
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.

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

Rebuild and execute the captured target call.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

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

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

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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

```python
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.

```python
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.

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

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

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