Patch TorchAO Float8Tensor conversion ops for ZeroGPU

zerogpu_tinker_service.py

@@ -25,6 +25,10 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
+import torch
+
+torch.__future__.set_swap_module_params_on_conversion(True)
+
 from fastapi import APIRouter, Depends, FastAPI, Header, HTTPException
 from pydantic import BaseModel
 
@@ -383,8 +387,8 @@ def _disable_deepgemm_for_zerogpu() -> None:
     _log_event("zerogpu_deepgemm_disabled")
 
 
-def _patch_torchao_float8_empty_like() -> None:
-    """Patch TorchAO Float8Tensor empty_like for ZeroGPU tensor packing."""
+def _patch_torchao_float8_tensor_ops() -> None:
+    """Patch TorchAO Float8Tensor ops used by ZeroGPU tensor packing."""
     if os.environ.get("HATCHERY_ZEROGPU_PATCH_TORCHAO_EMPTY_LIKE", "1").lower() in {
         "0",
         "false",
@@ -395,35 +399,26 @@ def _patch_torchao_float8_empty_like() -> None:
     try:
         import inspect
 
-        import torch
         from torch.utils._python_dispatch import return_and_correct_aliasing
         from torchao.quantization import Float8Tensor
     except Exception as exc:  # noqa: BLE001
         _log_event(
-            "zerogpu_torchao_empty_like_patch_skipped",
+            "zerogpu_torchao_float8_ops_patch_skipped",
             error_type=type(exc).__name__,
             error=str(exc),
         )
         return
 
-    if getattr(Float8Tensor, "_hatchery_empty_like_patch", False):
+    if getattr(Float8Tensor, "_hatchery_float8_ops_patch", False):
         return
 
     aten = torch.ops.aten
     implements = Float8Tensor.implements
+    init_params = inspect.signature(Float8Tensor).parameters
 
-    @implements(aten.empty_like.default)
-    def _(func, types, args, kwargs):
-        original_kwargs = dict(kwargs or {})
-        self = args[0]
-        out_dtype = original_kwargs.pop("dtype", None) or self.dtype
-        qdata = aten.empty_like.default(self.qdata, **original_kwargs)
-        scale_kwargs = {k: v for k, v in original_kwargs.items() if k != "dtype"}
-        scale = aten.empty_like.default(self.scale, **scale_kwargs)
-
-        init_params = inspect.signature(self.__class__).parameters
+    def _new_float8_tensor(self, qdata, scale, out_dtype):
         if "hp_value_lb" in init_params:
-            new = self.__class__(
+            return self.__class__(
                 qdata,
                 scale,
                 self.block_size,
@@ -434,20 +429,52 @@ def _patch_torchao_float8_empty_like() -> None:
                 self.kernel_preference,
                 out_dtype,
             )
-        else:
-            new = self.__class__(
-                qdata,
-                scale,
-                self.block_size,
-                self.mm_config,
-                self.act_quant_kwargs,
-                self.kernel_preference,
-                out_dtype,
-            )
+        return self.__class__(
+            qdata,
+            scale,
+            self.block_size,
+            self.mm_config,
+            self.act_quant_kwargs,
+            self.kernel_preference,
+            out_dtype,
+        )
+
+    @implements(aten.empty_like.default)
+    def _(func, types, args, kwargs):
+        original_kwargs = dict(kwargs or {})
+        self = args[0]
+        out_dtype = original_kwargs.pop("dtype", None) or self.dtype
+        qdata = aten.empty_like.default(self.qdata, **original_kwargs)
+        scale_kwargs = {k: v for k, v in original_kwargs.items() if k != "dtype"}
+        scale = aten.empty_like.default(self.scale, **scale_kwargs)
+        new = _new_float8_tensor(self, qdata, scale, out_dtype)
+        return return_and_correct_aliasing(func, args, kwargs, new)
+
+    @implements(aten._to_copy.default)
+    def _(func, types, args, kwargs):
+        kwargs = dict(kwargs or {})
+        self = args[0]
+        out_dtype = kwargs.get("dtype", None) or self.dtype
+        inner_kwargs = {k: v for k, v in kwargs.items() if k != "dtype"}
+        qdata = aten._to_copy.default(self.qdata, **inner_kwargs)
+        scale = aten._to_copy.default(self.scale, **inner_kwargs)
+        new = _new_float8_tensor(self, qdata, scale, out_dtype)
+        return return_and_correct_aliasing(func, args, kwargs, new)
+
+    @implements([aten.detach.default, aten.clone.default])
+    def _(func, types, args, kwargs):
+        kwargs = dict(kwargs or {})
+        self = args[0]
+        new = _new_float8_tensor(
+            self,
+            func(self.qdata, **kwargs),
+            func(self.scale, **kwargs),
+            self.dtype,
+        )
         return return_and_correct_aliasing(func, args, kwargs, new)
 
-    Float8Tensor._hatchery_empty_like_patch = True
-    _log_event("zerogpu_torchao_empty_like_patched")
+    Float8Tensor._hatchery_float8_ops_patch = True
+    _log_event("zerogpu_torchao_float8_ops_patched")
 
 
 def _ensure_supported_model(base_model: str, rank: int) -> None:
@@ -473,7 +500,7 @@ def _load_startup_model() -> Any:
     from transformers import AutoModelForCausalLM, AutoTokenizer
 
     _disable_deepgemm_for_zerogpu()
-    _patch_torchao_float8_empty_like()
+    _patch_torchao_float8_tensor_ops()
     base_model = _supported_model_id()
     started = time.monotonic()
     _log_event("zerogpu_startup_model_load_start", base_model=base_model, rank=_supported_rank())