aitune.torch.backend.torch_eager | NVIDIA AITune Documentation

Torch eager backend.

Module Contents

Classes

Name	Description
`TorchEagerBackend`	Backend that runs the model in eager mode with/without autocast.
`TorchEagerBackendConfig`	Configuration for torch eager backend.

API

class aitune.torch.backend.torch_eager.TorchEagerBackend(
    config: aitune.torch.backend.torch_eager.TorchEagerBackendConfig | None = None
)

Bases: Backend

Backend that runs the model in eager mode with/without autocast.

Note: inference is done with torch.no_grad() context. The torch.inference_mode() context must not be used as it would require outputs from a model to be used with same inference mode - this would be confusing to a user and required code changes from the user.

STATE_CONFIG

= 'config'

STATE_DEVICE

= 'device'

STATE_ORIG_MODULE

= 'orig_module'

STATE_OUTPUT_DTYPE

= 'output_dtype'

STATE_TYPE

= 'type'

_config

= config or TorchEagerBackendConfig()

aitune.torch.backend.torch_eager.TorchEagerBackend._activate()

Activates runner.

aitune.torch.backend.torch_eager.TorchEagerBackend._build(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    cache_dir: pathlib.Path
) -> aitune.torch.backend.backend.Backend

Builds the model.

aitune.torch.backend.torch_eager.TorchEagerBackend._deactivate()

Deactivates runner.

aitune.torch.backend.torch_eager.TorchEagerBackend._deploy()

Deploys the backend.

aitune.torch.backend.torch_eager.TorchEagerBackend._get_required_casting_dtype(
    module: torch.nn.Module,
    data: list[aitune.torch.module.recording_module.Sample]
) -> torch.dtype | None

Get the required casting dtype of the module by running a sample inference with and without autocast.

If the dtype of the output is different with and without autocast, return the dtype of the output without autocast. Otherwise, return None.

Parameters:

module

nn.Module

The module to get the dtype from.

data

list[Sample]

List of sample inputs to run through the module.

Returns: torch.dtype | None

torch.dtype: The required casting dtype. Returns None if no casting is required.

aitune.torch.backend.torch_eager.TorchEagerBackend._infer(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Runs inference with the given arguments. Does not use autocast.

It can be replaced at runtime by _infer_with_autocast.

Parameters:

*args

AnyDefaults to ()

inference arguments

**kwargs

AnyDefaults to {}

inference keyword arguments

Returns: Any

The result of the inference.

aitune.torch.backend.torch_eager.TorchEagerBackend._infer_with_autocast(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Runs inference with the given arguments.

Parameters:

*args

AnyDefaults to ()

inference arguments

**kwargs

AnyDefaults to {}

inference keyword arguments

Returns: Any

The result of the inference.

aitune.torch.backend.torch_eager.TorchEagerBackend._save_config(
    cache_dir: pathlib.Path
)

Store the backend configuration to a file.

aitune.torch.backend.torch_eager.TorchEagerBackend.describe() -> str

Returns the description of the backend.

aitune.torch.backend.torch_eager.TorchEagerBackend.from_dict(
    module: torch.nn.Module | None,
    state_dict: dict
)

classmethod

Creates a backend from a state_dict.

aitune.torch.backend.torch_eager.TorchEagerBackend.is_jit() -> bool

Returns True if the backend is a JIT backend.

aitune.torch.backend.torch_eager.TorchEagerBackend.key() -> str

Returns the key of the backend.

aitune.torch.backend.torch_eager.TorchEagerBackend.to_dict()

Returns the state_dict of the backend.

class aitune.torch.backend.torch_eager.TorchEagerBackendConfig(
    autocast_enabled: bool = False,
    autocast_dtype: torch.dtype | None = None
)

Dataclass

Bases: BackendConfig

Configuration for torch eager backend.

Parameters:

autocast_enabled

boolDefaults to False

If True, enable autocast.

autocast_dtype

torch.dtypeDefaults to None

The dtype to use for autocast.

autocast_dtype

dtype | None = None

autocast_enabled

bool = False