overrides

fastforward.overrides.DisableQuantizationOverride()

Override to disable quantization.

Attach DisableQuantizationOverride instance as quantizer override to disable quantization. Quantization is enabled/disabled using the enable_quantization and disable_quantization methods for all quantizers the instance is attached to.

Quantizers can be 'bulk' attached to using the attach_to and detach methods. Note that detach will only detach from quantizers which where attached to using the attach_to methods, i.e., if an instance of DisableQuantizationOverride is registered as override to a quantizers using different means, it must be detached separately.

Source code in fastforward/overrides.py

def __init__(self) -> None:
    self._quantization_enabled = False
    self._handles: list[override.OverrideHandle] = []

quantization_enabled property

True if quantization is enabled, False otherwise.

__call__(_context, callback, args, kwargs)

Override function for quantizer disabling.

Source code in fastforward/overrides.py

def __call__(
    self,
    _context: Any,
    callback: Callable[..., torch.Tensor],
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> torch.Tensor:
    """Override function for quantizer disabling."""
    if self._quantization_enabled:
        return callback(*args, **kwargs)
    else:
        return _extract_data_from_args(*args, **kwargs)

__repr__()

Source code in fastforward/overrides.py

@typing_override
def __repr__(self) -> str:
    return f"{type(self).__name__}(quantization_enabled={self._quantization_enabled})"

attach_to(quantizers)

Attach this override to one or more quantizers.

Parameters:

Name	Type	Description	Default
quantizers	Quantizer | QuantizerCollection | Iterable[Quantizer]	Either a single quantizer, a QuantizerCollection obtained using ff.find_quantizers or an iterable of Quantizers. Attach this override to each.	required

Source code in fastforward/overrides.py

def attach_to(
    self, quantizers: ff.nn.Quantizer | QuantizerCollection | Iterable[ff.nn.Quantizer]
) -> None:
    """Attach this override to one or more quantizers.

    Args:
        quantizers: Either a single quantizer, a `QuantizerCollection`
            obtained using `ff.find_quantizers` or an iterable of `Quantizer`s.
            Attach this override to each.
    """
    if isinstance(quantizers, ff.nn.Quantizer):
        self._handles.append(quantizers.register_override(self))
    elif isinstance(quantizers, QuantizerCollection):
        for quantizer in quantizers.modules():
            assert isinstance(quantizer, ff.nn.Quantizer)
            self.attach_to(quantizer)
    else:
        for quantizer in quantizers:  # type: ignore[union-attr]
            self.attach_to(quantizer)

detach()

Detach this override.

Detach from all quantizers it was attached to using the attach_to method.

Source code in fastforward/overrides.py

def detach(self) -> None:
    """Detach this override.

    Detach from all quantizers it was attached to using the
    `attach_to` method.
    """
    for handle in self._handles:
        handle.remove()
    self._handles = []

disable_quantization()

Disable quantization.

See the docstring of enable_quantization for more information.

Source code in fastforward/overrides.py

def disable_quantization(self) -> "_QuantizationContext":
    """Disable quantization.

    See the docstring of `enable_quantization` for more information.
    """
    return self.enable_quantization(enabled=False)

enable_quantization(enabled=True)

Enable quantization.

More specifically, this instance will not disable quantization for any quantizers it is attached to. Other instance, or other methods may still disable quantization.

Parameters:

Name	Type	Description	Default
enabled	bool	True if quantization must be enabled, False if it must be disabled.	True

Source code in fastforward/overrides.py

def enable_quantization(self, enabled: bool = True) -> "_QuantizationContext":
    """Enable quantization.

    More specifically, this instance will not disable quantization for any
    quantizers it is attached to. Other instance, or other methods may
    still disable quantization.

    Args:
        enabled: True if quantization must be enabled, False if it must be disabled.
    """
    context = _QuantizationContext(self, self._quantization_enabled)
    self._quantization_enabled = enabled
    return context

fastforward.overrides.disable_quantization(model)

Disable quantization for all quantizers in model within context.

The global strict_quantization flag is also set to False during the context.

Parameters:

Name	Type	Description	Default
model	Module	The model for which all quantizers are disabled.	required

Source code in fastforward/overrides.py

@contextlib.contextmanager
def disable_quantization(model: torch.nn.Module) -> Generator[None, None, None]:
    """Disable quantization for all quantizers in `model` within context.

    The global `strict_quantization` flag is also set to `False` during the context.

    Args:
        model: The model for which all quantizers are disabled.
    """
    handles: list[override.OverrideHandle] = []
    for _, quantizer in named_quantizers(model):
        handles.append(quantizer.register_override(DisableQuantizationOverride()))

    try:
        with ff.strict_quantization(False):
            yield
    finally:
        for handle in handles:
            handle.remove()

fastforward.overrides.enable_quantization(model)

Enable quantization for all quantizers in model within context.

Note that this context manager does not change the strict_quantization flag. To also (temporarily) change the strict_quantization flag use fastforward.quantization.strict_quantization.strict_quantization_for_module

Parameters:

Name	Type	Description	Default
model	Module	The model for which all quantizers are enabled.	required

Source code in fastforward/overrides.py

@contextlib.contextmanager
def enable_quantization(model: torch.nn.Module) -> Generator[None, None, None]:
    """Enable quantization for all quantizers in `model` within context.

    Note that this context manager does not change the `strict_quantization` flag.
    To also (temporarily) change the `strict_quantization` flag use
    `fastforward.quantization.strict_quantization.strict_quantization_for_module`

    Args:
        model: The model for which all quantizers are enabled.
    """
    with contextlib.ExitStack() as exit_stack:
        for _, quantizer in named_quantizers(model):
            for quantizer_override in quantizer.overrides:
                if isinstance(quantizer_override, DisableQuantizationOverride):
                    exit_stack.enter_context(quantizer_override.enable_quantization())
        yield