nemo_automodel.autonvtx | NVIDIA NeMo AutoModel

Module Contents

Name	Description
`_add_nvtx_hooks`	Add NVTX range hooks to a model’s forward and optionally backward passes.
`_get_active_ranges`	Get the set of currently active NVTX ranges for this thread.
`patch`	Recursively patch a model with NVTX profiling annotations.

nemo_automodel.autonvtx._add_nvtx_hooks(
    model,
    name,
    add_backward_hooks = True
)

Add NVTX range hooks to a model’s forward and optionally backward passes.

nemo_automodel.autonvtx._get_active_ranges()

Get the set of currently active NVTX ranges for this thread.

nemo_automodel.autonvtx.patch(
    model,
    name = None,
    add_backward_hooks = True
)

Recursively patch a model with NVTX profiling annotations.

Prevents duplicate scopes when activation checkpointing reruns forward passes.

nemo_automodel.autonvtx.__all__ = ['patch']

nemo_automodel.autonvtx._thread_local = local()