aitune.torch.module.recording_module | NVIDIA AITune Documentation

Inplace wrap modules.

Module Contents

Classes

Name	Description
`RecordingModule`	Module that records samples for tuning.

Data

INPUT_METADATA_PREFIX

OUTPUT_METADATA_PREFIX

logger

API

class aitune.torch.module.recording_module.RecordingModule(
    module: torch.nn.Module,
    name: str,
    config: aitune.torch.config.AITuneConfig | None = None,
    dynamic_shapes: aitune.torch.dynamic_shapes.DynamicShapes | None = None,
    cache_dir_resolver: collections.abc.Callable[[], pathlib.Path]
)

Module that records samples for tuning.

_config

= config if config is not None else global_config

_dynamic_shapes

= dynamic_shapes or {}

_forward_call

= module.__call__

_forward_signature

= ForwardSignature.from_callable(module.forward)

_graph_specs

OrderedDict[SampleMetadata, GraphSpec] = OrderedDict()

_graphs_counter

= itertools.count()

_samples

dict[SampleMetadata, SampleStore] = {}

device

Get the device of the module.

forward_signature

ForwardSignature

Get the forward signature used to normalize inputs.

graph_specs

list[GraphSpec]

Get the graph specs.

is_ready_for_optimization

bool

Check if the module is ready for tuning.

All graphs would have proper amount of samples.

aitune.torch.module.recording_module.RecordingModule.__call__(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Record a sample and run the module.

Before calling forward, we need to make input sample metadata and a copy of args and kwargs, since the forward call may have side effects on the inputs (like KV cache in LLM models).

aitune.torch.module.recording_module.RecordingModule._copy_inputs(
    args: tuple,
    kwargs: dict
) -> tuple[tuple, dict]

staticmethod

Create deep copies of args and kwargs.

Parameters:

args

tuple

Positional arguments tuple.

kwargs

dict

Keyword arguments dict.

Returns: tuple[tuple, dict]

Tuple containing deep copies of args and kwargs.

Raises:

RuntimeError: If cannot copy model inputs with tip to fix it

aitune.torch.module.recording_module.RecordingModule._resolve_and_validate_dynamic_shapes(
    inputs_metadata: aitune.torch.module.sample_metadata.SampleMetadata
) -> aitune.torch.dynamic_shapes.DynamicShapes

Match explicit shape definitions to tensors in a recorded sample.

The returned mapping is a validated, graph-specific subset of self._dynamic_shapes containing only definitions whose paths resolve to tensor inputs in inputs_metadata. Unconfigured tensor inputs remain represented by their TensorSpec and continue using shapes inferred from recorded samples.

aitune.torch.module.recording_module.RecordingModule.record_sample(
    inputs,
    inputs_metadata,
    outputs_metadata,
    dynamic_shapes: aitune.torch.dynamic_shapes.DynamicShapes | None = None,
    post_inputs_metadata: aitune.torch.module.sample_metadata.SampleMetadata | None = None
) -> None

Record a sample from the module.

aitune.torch.module.recording_module.RecordingModule.samples_for_graph_spec(
    graph_spec: aitune.torch.module.graph_spec.GraphSpec
) -> aitune.torch.module.sample_store.SampleStore

Get the samples.

aitune.torch.module.recording_module.RecordingModule.validate_dynamic_shape_paths_recorded() -> None

Validate that every configured input path matched a tensor during recording.

aitune.torch.module.recording_module.INPUT_METADATA_PREFIX = 'input'

aitune.torch.module.recording_module.OUTPUT_METADATA_PREFIX = 'output'

aitune.torch.module.recording_module.logger = logging.getLogger(__name__)