nemo_automodel.autonvtx

Package Contents

Functions

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

Data

_thread_local
__all__

API

nemo_automodel.autonvtx._thread_local

‘local(…)’

nemo_automodel.autonvtx._get_active_ranges()

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

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.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’]