Skip to content

quantizer_annotations

Classes#

fastforward.quantization.quantizer_annotations.TraceTensor(data) #

Bases: Tensor


              flowchart TD
              fastforward.quantization.quantizer_annotations.TraceTensor[TraceTensor]

              

              click fastforward.quantization.quantizer_annotations.TraceTensor href "" "fastforward.quantization.quantizer_annotations.TraceTensor"

A torch.Tensor subclass that tracks the preceding operation and quantizer.

TraceTensor instances store a reference to the previous_function (the operation that produced this tensor, set via __torch_function__) and the previous_quantizer (the quantizer layer that consumed or produced this tensor, set via quantizer_override). This tracking facilitates the annotation of Quantizer layers with Tag(before/after:func_name) metadata, allowing the user to search specific layers in the graph.

[source]

Attributes#
previous_function = None instance-attribute #
previous_quantizer = None instance-attribute #

Functions#

fastforward.quantization.quantizer_annotations.annotate_operator_metadata(model, sample_input) #

Trace model execution and annotate Quantizer metadata with preceding and succeeding operators.

This function performs a lightweight tracing of the model's forward pass using a special tracing tensor (TraceTensor) and quantizer overrides. It identifies the operators that immediately precede and follow each ff.nn.Quantizer instance in the execution flow and adds this information as Tags to the metadata (quant_metadata) within the Quantizer objects.

The annotation happens in-place on the Quantizer objects within the provided model. We assume that ff.quantize has been called on the model to ensure that the model contains Quantizer instances that need to be annotated.

Parameters:

Name Type Description Default
model QuantizedModule

The torch.nn.Module instance to be traced and annotated. It is assumed that ff.quantize has been called on this model.

 required
sample_input Tensor

A sample input tensor to the model. This is used to drive the forward pass for tracing purposes. If the module has conditional logic based on input data, ensure that sample_input covers all possible branches or that you run this function with multiple sample inputs to cover all branches.

 required

[source]

fastforward.quantization.quantizer_annotations.quantizer_override(quantizer, callback, args, kwargs) #

Register the operations that precede and follow this quantizer forward pass as Tags.

Combined with TraceTensor, this function will add both the preceeding and following operators to each quantizer layer as Tag objects: (Tag("before") / func.__name__)andTag("after") / func.name`).

Parameters:

Name Type Description Default
quantizer Quantizer

A quantizer layer.

 required
callback Callable[..., Any]

The quantizer layer forward call.

 required
args tuple[Any, ...]

The arguments passed to the quantizer layer forward call.

 required
kwargs dict[str, Any]

Any key-word arguments passed to the quantizer layer forward call.

 required

[source]