nemo_automodel.components.cuda_graphs

CUDA graph components.

Submodules

nemo_automodel.components.cuda_graphs.partial

Package Contents

Classes

Name	Description
`PartialCudaGraphManager`	Capture selected TE attention and MoE submodules after one eager iteration.

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.