Repository: foundation-model-stack/fms-model-optimizer
Pull request #96: feat: Update syntax of custom torch ops
Author: login: andrea-fasoli
Labels: name: feat
Created: 2025-04-14T23:17:11Z | Merged: 2025-04-23T13:00:08Z

### Description of the change

This PR updates the syntax of registration of a custom op with pytorch. The purpose of registration is to create a custom operation that can be inserted as custom node in the computational graph without inducing a graph break.

From PyTorch 2.4, new `torch.library` functions have been introduced: 
- `custom_op` replacing `impl`
- `register_fake` replacing `impl_abstract`
They streamline earlier implementations of the custom op registration process. With the new syntax, there is no need for an additional op definition using `torch.library.define`. 

The earlier syntax is deprecated from PyTorch >= 2.6: `impl_abstract` redirects to `register_fake` and throws a warning.
However, the new functions introduce a severe lower bound to `fms-mo`-supported PyTorch versions (>= 2.4), so we may want to hold to this update until later on.

### Was the PR tested

- [X] I have ensured all unit tests pass


Diff:
diff --git a/fms_mo/aiu_addons/gptq/gptq_aiu_op.py b/fms_mo/aiu_addons/gptq/gptq_aiu_op.py
index d958f38e..b9199f6b 100644
--- a/fms_mo/aiu_addons/gptq/gptq_aiu_op.py
+++ b/fms_mo/aiu_addons/gptq/gptq_aiu_op.py
@@ -17,6 +17,7 @@
 import logging
 
 # Third Party
+from packaging.version import Version
 import torch
 
 # pylint: disable=unused-argument
@@ -25,6 +26,36 @@
 logger = logging.getLogger(__name__)
 
 
+def implement_op_decorator(op_namespace_id):
+    """Version-dependent decorator for custom op implementation.
+    Always compare against pytorch version in current environment.
+    """
+
+    torch_version = Version(torch.__version__.split("+", maxsplit=1)[0])
+
+    def decorator(func):
+        if torch_version < Version("2.4"):
+            return torch.library.impl(op_namespace_id, "default")(func)
+        return torch.library.custom_op(op_namespace_id, mutates_args=())(func)
+
+    return decorator
+
+
+def register_op_decorator(op_namespace_id):
+    """Version-dependent decorator for custom op registration.
+    Always compare against pytorch version in current environment.
+    """
+
+    torch_version = Version(torch.__version__.split("+", maxsplit=1)[0])
+
+    def decorator(func):
+        if torch_version < Version("2.4"):
+            return torch.library.impl_abstract(op_namespace_id)(func)
+        return torch.library.register_fake(op_namespace_id)(func)
+
+    return decorator
+
+
 def register_aiu_gptq_op():
     """Register AIU-specific op to enable torch compile without graph break.
     The op preserves I/O shapes of a `X @ W^T` matmul but performs no operation.
@@ -36,17 +67,33 @@ def register_aiu_gptq_op():
     ):
         logger.warning("AIU op has already been registered")
         return
-
     op_namespace_id = "gptq_gemm::i4f16_fxinputs_aiu"
-    torch.library.define(
-        op_namespace_id,
-        "(Tensor x, Tensor qw, Tensor qzeros, Tensor scales, Tensor g_idx) -> Tensor",
-    )
+    if Version(torch.__version__.split("+", maxsplit=1)[0]) < Version("2.4"):
+        torch.library.define(
+            op_namespace_id,
+            "(Tensor x, Tensor qw, Tensor qzeros, "
+            "Tensor scales, Tensor g_idx) -> Tensor",
+        )
 
     # Add implementations for the operator
-    @torch.library.impl(op_namespace_id, "default")
-    def i4f16_fxinputs_aiu(x, qw, qzeros, scales, g_idx):
-        # on AIU, GPTQ qw is [out_feat, in_feat]
+    @implement_op_decorator(op_namespace_id)
+    def i4f16_fxinputs_aiu(
+        x: torch.Tensor,
+        qw: torch.Tensor,
+        qzeros: torch.Tensor,
+        scales: torch.Tensor,
+        g_idx: torch.Tensor,
+    ) -> torch.Tensor:
+        """Implement fake processing of GPTQ W4A16 matmul. The purpose is to create a
+        node on the computational graph to be captured during compiling for AIU.
+
+        Instead of computing the weight decompression and matmul, this function returns
+        a zero tensor with the expected shape.
+
+        NOTE: on AIU, GPTQ qw is [out_feat, in_feat], while AutoGPTQ saves the quantized
+        weights as [in_feat, out_feat]
+        """
+
         outshape = x.shape[:-1] + (qw.shape[0],)
         x = x.view(-1, x.shape[-1])
         output = torch.zeros(
@@ -56,8 +103,10 @@ def i4f16_fxinputs_aiu(x, qw, qzeros, scales, g_idx):
         )
         return output.view(outshape)
 
-    @torch.library.impl_abstract(op_namespace_id)
-    def i4f16_fxinputs_aiu_abstract(x, qw, qzeros, scales, g_idx):
+    @register_op_decorator(op_namespace_id)
+    def _(x, qw, qzeros, scales, g_idx):
+        """OP template of I/O sizes"""
+
         outshape = x.shape[:-1] + (qw.shape[0],)
         return torch.empty(
             outshape,
diff --git a/fms_mo/aiu_addons/i8i8/i8i8_aiu_op.py b/fms_mo/aiu_addons/i8i8/i8i8_aiu_op.py
index 0d38836c..41aa896f 100644
--- a/fms_mo/aiu_addons/i8i8/i8i8_aiu_op.py
+++ b/fms_mo/aiu_addons/i8i8/i8i8_aiu_op.py
@@ -17,11 +17,10 @@
 import logging
 
 # Third Party
+from packaging.version import Version
 import torch
 import torch.nn.functional as F
 
-logger = logging.getLogger(__name__)
-
 # pylint: disable=unused-argument
 # i8i8 op must be registered with specific I/O, even if not in use by the op function
 
@@ -29,6 +28,38 @@
 # torch.nn.functional.linear not recognized as callable
 # open issue in PyLint: https://github.com/pytorch/pytorch/issues/119482
 
+logger = logging.getLogger(__name__)
+
+
+def implement_op_decorator(op_namespace_id):
+    """Version-dependent decorator for custom op implementation.
+    Always compare against pytorch version in current environment.
+    """
+
+    torch_version = Version(torch.__version__.split("+", maxsplit=1)[0])
+
+    def decorator(func):
+        if torch_version < Version("2.4"):
+            return torch.library.impl(op_namespace_id, "default")(func)
+        return torch.library.custom_op(op_namespace_id, mutates_args=())(func)
+
+    return decorator
+
+
+def register_op_decorator(op_namespace_id):
+    """Version-dependent decorator for custom op registration.
+    Always compare against pytorch version in current environment.
+    """
+
+    torch_version = Version(torch.__version__.split("+", maxsplit=1)[0])
+
+    def decorator(func):
+        if torch_version < Version("2.4"):
+            return torch.library.impl_abstract(op_namespace_id)(func)
+        return torch.library.register_fake(op_namespace_id)(func)
+
+    return decorator
+
 
 def register_aiu_i8i8_op():
     """Register AIU-specific op to enable torch compile without graph break.
@@ -41,26 +72,26 @@ def register_aiu_i8i8_op():
     if hasattr(torch.ops, "fms_mo") and hasattr(torch.ops.fms_mo, "i8i8_aiu"):
         logger.warning("AIU op has already been registered")
         return
-
     op_namespace_id = "fms_mo::i8i8_aiu"
-    torch.library.define(
-        op_namespace_id,
-        "(Tensor x, Tensor weight, Tensor bias, Tensor qdata, "
-        "str weight_quant_type, str activ_quant_type, "
-        "bool smoothquant) "
-        "-> Tensor",
-    )
+    if Version(torch.__version__.split("+", maxsplit=1)[0]) < Version("2.4"):
+        torch.library.define(
+            op_namespace_id,
+            "(Tensor x, Tensor weight, Tensor bias, Tensor qdata, "
+            "str weight_quant_type, str activ_quant_type, "
+            "bool smoothquant) "
+            "-> Tensor",
+        )
 
-    @torch.library.impl(op_namespace_id, "default")
+    @implement_op_decorator(op_namespace_id)
     def i8i8_aiu(
-        x,
-        weight,
-        bias,
-        qdata,
-        weight_quant_type,
-        activ_quant_type,
-        smoothquant,
-    ):
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        bias: torch.Tensor,
+        qdata: torch.Tensor,
+        weight_quant_type: str,
+        activ_quant_type: str,
+        smoothquant: bool,
+    ) -> torch.Tensor:
         """Implement addmm of X and W.
         Support various quantization options for weights and activations.
 
@@ -86,16 +117,8 @@ def i8i8_aiu(
 
         return F.linear(x_dq.to(dtype), w_dq.to(dtype), bias.to(dtype))
 
-    @torch.library.impl_abstract(op_namespace_id)
-    def i8i8_aiu_abstract(
-        x,
-        weight,
-        bias,
-        qdata,
-        weight_quant_type,
-        activ_quant_type,
-        smoothquant,
-    ):
+    @register_op_decorator(op_namespace_id)
+    def _(x, weight, bias, qdata, weight_quant_type, activ_quant_type, smoothquant):
         """OP template of I/O sizes"""
 
         outshape = x.size()[:-1] + (weight.size(0),)


Reviews:
user: login: chichun-charlie-liu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-21T13:56:55Z
user: login: andrea-fasoli
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-21T16:50:57Z
user: login: chichun-charlie-liu
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-23T12:59:55Z

Review comments:
diff_hunk: @@ -17,12 +17,36 @@
 import logging
 
 # Third Party
+from packaging.version import Version
 import torch
 
 # pylint: disable=unused-argument
 # gptq op must be registered with specific I/O, even if not in use by the op function
 
 logger = logging.getLogger(__name__)
+torch_version = Version(torch.__version__.split("+", maxsplit=1)[0])
+
+
+def implement_op_decorator(pt_ver, op_namespace_id):
path: fms_mo/aiu_addons/gptq/gptq_aiu_op.py
user: login: chichun-charlie-liu
body: do we really need to pass `pt_ver` as an arg to this func? it can access `torch_version` defined on L27 directly. Unless there is a case we want to register using a syntax that is lower than the current installed PT version? 
created_at: 2025-04-21T13:56:54Z
author_association: COLLABORATOR
diff_hunk: @@ -17,12 +17,36 @@
 import logging
 
 # Third Party
+from packaging.version import Version
 import torch
 
 # pylint: disable=unused-argument
 # gptq op must be registered with specific I/O, even if not in use by the op function
 
 logger = logging.getLogger(__name__)
+torch_version = Version(torch.__version__.split("+", maxsplit=1)[0])
+
+
+def implement_op_decorator(pt_ver, op_namespace_id):
path: fms_mo/aiu_addons/gptq/gptq_aiu_op.py
user: login: andrea-fasoli
body: `pt_ver` is not needed at this time but l plan to move these decorators somewhere under `utils` and import them (they are shared between gptq and int8), so I'd prefer for them to be more general and not force users to declare a global `torch_version` variable, even in future addons.
created_at: 2025-04-21T16:50:57Z
author_association: COLLABORATOR

PR comments:
user: login: chichun-charlie-liu
created_at: 2025-04-18T15:19:59Z
body: can we use a "version check" approach like the one we used for external kernels [here](https://github.com/foundation-model-stack/fms-model-optimizer/blob/main/fms_mo/custom_ext_kernels/utils.py#L52)?

author_association: COLLABORATOR
user: login: andrea-fasoli
created_at: 2025-04-18T21:42:06Z
body: addons updated with branching based on pytorch version.
There's some duplication of the version-dependent decorators between GPTQ and INT8. We can eventually move these to `fms_mo/utils/aiu_utils.py` and import them from there.
Notice FMS forces `torch==2.5.1` at this time, so this version check is only relevant when FMS-MO is used independently from FMS.
author_association: COLLABORATOR