how to implement a customer precision with torchao

[qmpzzpmq]

Bug description

hi, I implemented a customer precision with torchao.quantization.Int8WeightOnlyConfig, but unfortunately it looks like customer precision plugin integrated is not allowed in the lightning since the function _convert_precision_to_unified_args in lightning/fabric/connector.py check only the offical implement precision. Does anyone know, how can I implement it?

What version are you seeing the problem on?

v2.5

Reproduced in studio

No response

How to reproduce the bug

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):

More info

No response

cc @ethanwharris

[SkafteNicki]

Hi @qmpzzpmq, thanks for raising this issue :)
Can you provide a bit more info on where in the process you are running into problems? It should be as simple as creating a custom precision plugin from the Precision class and then passing that to trainer:

import torch
import torch.nn as nn
from lightning import LightningModule, Trainer
from lightning.pytorch.plugins import Precision


class CustomInt8Precision(Precision):   # dummy example
    def __init__(self):
        self.precision = "custom-int8"

    def convert_input(self, data):
        return data

    def convert_output(self, data):
        return data

    def convert_module(self, module):
        return module

    def convert_optimizer(self, optimizer):
        return optimizer

    def teardown(self):
        pass


class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(10, 1)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = nn.functional.mse_loss(logits, y)
        return loss

custom_precision = CustomInt8Precision()
trainer = Trainer(plugins=[custom_precision], max_epochs=1)

[qmpzzpmq]

# https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/plugins/precision/bitsandbytes.py
from contextlib import AbstractContextManager
from typing import Any, Literal

import torch
import torch.nn as nn
from lightning.fabric.plugins.precision.precision import Precision
from lightning.fabric.plugins.precision.utils import (
    _convert_fp_tensor,
    _DtypeContextManager,
)
from lightning_utilities.core.apply_func import apply_to_collection

# TorchAO quantization imports
from torchao.quantization import (
    Int4WeightOnlyConfig,
    Int8WeightOnlyConfig,
    quantize_,
)
from typeguard import typechecked
from typing_extensions import override


# https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.plugins.precision.Precision.html#lightning.pytorch.plugins.precision.Precision
class AFP16WIntPrecision(Precision):
    @typechecked
    def __init__(
        self,
        activation_dtype: Literal["fp16", "bf16"] = "fp16",
        weight_dtype: Literal["int4", "int8"] = "int8",
        group_size: int = 32,
    ) -> None:
        super().__init__()

        assert activation_dtype in [
            "fp16",
            "bf16",
        ], "Unsupported activation dtype"
        assert weight_dtype in ["int4", "int8"], "Unsupported weight dtype"

        self.activation_dtype_name = activation_dtype
        self.activation_dtype = (
            torch.float16 if activation_dtype == "fp16" else torch.bfloat16
        )
        self.weight_dtype = weight_dtype
        self.group_size = group_size
        if self.weight_dtype == "int4":
            raise NotImplementedError(
                "int4 quantization is not supported yet."
            )

    def _quantize_weights(self, model: nn.Module) -> None:
        """
        Quantize model weights to int4/int8
        while keeping activations in fp16/bf16.
        """
        if self.weight_dtype == "int4":
            # Use TorchAO's Int4WeightOnlyConfig for 4-bit quantization
            quantize_(model, Int4WeightOnlyConfig(group_size=self.group_size))
        elif self.weight_dtype == "int8":
            # Use TorchAO's Int8WeightOnlyConfig for 8-bit quantization
            quantize_(model, Int8WeightOnlyConfig(group_size=self.group_size))
        else:
            raise ValueError(f"Unsupported weight dtype: {self.weight_dtype}")
        return model

    @override
    def tensor_init_context(self) -> AbstractContextManager:
        return _DtypeContextManager(self.activation_dtype)

    @override
    def module_init_context(self) -> AbstractContextManager:
        return self.tensor_init_context()

    @override
    def convert_module(self, convert_module: nn.Module) -> nn.Module:
        """
        Convert module to the specified activation dtype.
        """
        return self._quantize_weights(convert_module)

    def forward_context(self) -> Any:
        """
        Return a context manager to set the correct dtype
        during forward pass."""
        return _DtypeContextManager(self.activation_dtype)

    @override
    def convert_input(self, data: Any) -> Any:
        return apply_to_collection(
            data,
            function=_convert_fp_tensor,
            dtype=torch.Tensor,
            dst_type=self.activation_dtype,
        )

    @override
    def convert_output(self, data: Any) -> Any:
        return apply_to_collection(
            data,
            function=_convert_fp_tensor,
            dtype=torch.Tensor,
            dst_type=torch.get_default_dtype(),
        )

    def backward(
        self, tensor: torch.Tensor, model: nn.Module, *args: Any, **kwargs: Any
    ) -> None:
        """Not used for inference."""
        raise ValueError(
            "Backward pass is not implemented for fp16*int8 inference"
            " precision."
        )

    def optimizer_step(
        self, optimizer: torch.optim.Optimizer, **kwargs: Any
    ) -> Any:
        """Not used for inference."""
        raise ValueError(
            "Optimizer step is not implemented for fp16*int8 inference"
            " precision."
        )

hi, @SkafteNicki here is my implement