> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/cuda_graphs/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

CUDA graph components.

## Submodules

* **[`nemo_automodel.components.cuda_graphs.partial`](/nemo-automodel/nemo_automodel/components/cuda_graphs/partial)**

## Package Contents

### Classes

| Name                                                                                                | Description                                                                 |
| --------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------- |
| [`PartialCudaGraphManager`](#nemo_automodel-components-cuda_graphs-partial-PartialCudaGraphManager) | Capture selected TE attention and MoE submodules after one eager iteration. |

### 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.

## Docs

- [nemo_automodel.components.cuda_graphs.partial](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/cuda_graphs/partial.md)