tracer

Tracing module for building a set of partial Graphs describing the model.

Classes

RobustTracer

A robust graph tracer for tracing dependencies between layers.

GraphCollection

A collection of fx graphs and tracing results.

Functions

recursive_trace

Recursively trace the model and all failing sub-modules and return graph collection.
class GraphCollection

Bases: object

A collection of fx graphs and tracing results.

__init__()

Initialize the graph collection.

Return type:

None

failure_msg(m)

Get the failure message of the module if it failed.

Parameters:

m (Module | Callable) –

Return type:

str | None

is_failed(m)

Check if node target failed during tracing.

Parameters:

m (Module | Callable) –

Return type:

bool

is_unvisited(m)

Check if the node target is unvisited.

Parameters:

m (Module | Callable) –

Return type:

bool

recursive_trace(model, concrete_args=None)

Recursively trace the model and all failing sub-modules and store graph collection.

Parameters:

  • model (Module) – The model to be traced.

  • concrete_args (Dict[str, Any] | None) – [EXPERIMENTAL; not backwards compatible!] The concrete arguments to be used for tracing of the root module (i.e. the model’s forward method). Note that if/when the tracing recurses, no concrete args are passed on! Defaults to None.

Return type:

None

Note that the graph collection contains the graph of the root module as well as all failing submodules. For failing modules, the graph will contain the failing module as a leaf and then separate graphs for each submodule within the failing module to maximize the coverage of the tracing results.

class RobustTracer

Bases: object

A robust graph tracer for tracing dependencies between layers.

Note that this tracer wraps unsupported ops into its parent module as a whole.

__init__()

Initialize the tracer.

failure_msg(m)

Get the failure message of the module.

Parameters:

m (Module | Callable) –

Return type:

str | None

is_failed(m)

Check if the module is failed.

Parameters:

m (Module | Callable) –

Return type:

bool

is_leaf_module(m, module_qualified_name)

Check if the given module is a leaf module.

Parameters:

  • m (Module) –

  • module_qualified_name (str) –

Return type:

bool

classmethod is_registered_leaf(m)

Check if module is registered as extra leaf module.

Parameters:

m (Module | Type[Module]) –

Return type:

bool

is_unvisited(m)

Check if the node target is unvisited.

Parameters:

m (Module | Callable) –

Return type:

bool

record_call_module(m, e=None)

Record the call_module operation.

Parameters:

  • m (Module) –

  • e (Exception | None) –

Return type:

None

classmethod register_leaf(m)

Register a module as extra leaf module.

Parameters:

m (Type[Module]) –

Return type:

None

trace(root, concrete_args=None)

Trace root using the fx tracer.

Unlike the standard tracer in fx, this tracer adds the following functionalities:

  1. Only torch.Tensors in concrete_args are converted to Proxy objects.

  2. The tracer doesn’t fail when encountering unsupported operations. Instead it will wrap the module containing the unsupported operation into a leaf node and repeat the tracing.

  3. The tracer doesn’t trace into leaf modules even if the leaf module is the root.

Parameters:

  • root (Module) –

  • concrete_args (Dict[str, Any] | None) –

Return type:

Graph

classmethod unregister_leaf(m)

Unregister a module as extra leaf module.

Parameters:

m (Type[Module]) –

Return type:

None

recursive_trace(model, concrete_args=None)

Recursively trace the model and all failing sub-modules and return graph collection.

Parameters:

  • model (Module) – The model to be traced.

  • concrete_args (Dict[str, Any] | None) – [EXPERIMENTAL; not backwards compatible!] The concrete arguments to be used for tracing of the root module (i.e. the model’s forward method). Note that if/when the tracing recurses, no concrete args are passed on! Defaults to None.

Returns:

The collection of fx graphs resulting from the tracing.

Return type:

GraphCollection

Note that the graph collection contains the graph of the root module as well as all failing submodules. For failing modules, the graph will contain the failing module as a leaf and then separate graphs for each submodule within the failing module to maximize the coverage of the tracing results.