Incorporate torch version acquisition into decorator

Signed-off-by: Andrea Fasoli <[email protected]>

Author: andrea-fasoli

fms_mo/aiu_addons/gptq/gptq_aiu_op.py

@@ -24,25 +24,32 @@
 # 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):
-    """Version-dependent decorator for custom op implementation."""
+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 pt_ver < Version("2.4"):
+        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(pt_ver, op_namespace_id):
-    """Version-dependent decorator for custom op registration."""
+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 pt_ver < Version("2.4"):
+        if torch_version < Version("2.4"):
             return torch.library.impl_abstract(op_namespace_id)(func)
         return torch.library.register_fake(op_namespace_id)(func)
 
@@ -61,15 +68,15 @@ def register_aiu_gptq_op():
         logger.warning("AIU op has already been registered")
         return
     op_namespace_id = "gptq_gemm::i4f16_fxinputs_aiu"
-    if torch_version < Version("2.4"):
+    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
-    @implement_op_decorator(torch_version, op_namespace_id)
+    @implement_op_decorator(op_namespace_id)
     def i4f16_fxinputs_aiu(
         x: torch.Tensor,
         qw: torch.Tensor,
@@ -96,7 +103,7 @@ def i4f16_fxinputs_aiu(
         )
         return output.view(outshape)
 
-    @register_op_decorator(torch_version, op_namespace_id)
+    @register_op_decorator(op_namespace_id)
     def _(x, qw, qzeros, scales, g_idx):
         """OP template of I/O sizes"""
 

fms_mo/aiu_addons/i8i8/i8i8_aiu_op.py

@@ -29,25 +29,32 @@
 # open issue in PyLint: https://github.com/pytorch/pytorch/issues/119482
 
 logger = logging.getLogger(__name__)
-torch_version = Version(torch.__version__.split("+", maxsplit=1)[0])
 
 
-def implement_op_decorator(pt_ver, op_namespace_id):
-    """Version-dependent decorator for custom op implementation."""
+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 pt_ver < Version("2.4"):
+        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(pt_ver, op_namespace_id):
-    """Version-dependent decorator for custom op registration."""
+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 pt_ver < Version("2.4"):
+        if torch_version < Version("2.4"):
             return torch.library.impl_abstract(op_namespace_id)(func)
         return torch.library.register_fake(op_namespace_id)(func)
 
@@ -66,7 +73,7 @@ def register_aiu_i8i8_op():
         logger.warning("AIU op has already been registered")
         return
     op_namespace_id = "fms_mo::i8i8_aiu"
-    if torch_version < Version("2.4"):
+    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, "
@@ -75,7 +82,7 @@ def register_aiu_i8i8_op():
             "-> Tensor",
         )
 
-    @implement_op_decorator(torch_version, op_namespace_id)
+    @implement_op_decorator(op_namespace_id)
     def i8i8_aiu(
         x: torch.Tensor,
         weight: torch.Tensor,
@@ -110,7 +117,7 @@ def i8i8_aiu(
 
         return F.linear(x_dq.to(dtype), w_dq.to(dtype), bias.to(dtype))
 
-    @register_op_decorator(torch_version, op_namespace_id)
+    @register_op_decorator(op_namespace_id)
     def _(x, weight, bias, qdata, weight_quant_type, activ_quant_type, smoothquant):
         """OP template of I/O sizes"""