fix: patch torchao dtype imports for unsloth

Allow the training Space to import Unsloth with stable CUDA PyTorch by shimming optional TorchAO dtype probes before Transformers loads.

Made-with: Cursor

Dockerfile.train

@@ -4,7 +4,7 @@ WORKDIR /app
 
 RUN apt-get update && apt-get install -y git build-essential && rm -rf /var/lib/apt/lists/*
 
-# Stable cu121 torch — register_constant shim applied at runtime for torchao compat
+# Stable cu121 torch — TorchAO import shims are applied at runtime before Unsloth loads.
 RUN pip install --no-cache-dir \
     torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 && \
     pip install --no-cache-dir \

training/space_runner.py

@@ -57,13 +57,14 @@ def _attach_log_capture():
         lg.setLevel(logging.INFO)
 
 
-def _run_training():
-    global _training_status
-    _training_status = "running"
-    _append_log("Thread started — patching torch/_pytree then importing unsloth...")
+def _patch_torchao_import_compat():
+    """Patch TorchAO import probes that expect newer/nightly torch symbols.
+
+    Training uses bitsandbytes 4-bit loading through Unsloth, not TorchAO
+    quantization. These aliases are only to let optional TorchAO modules import.
+    """
+    import torch
 
-    # torchao uses register_constant which only exists in torch nightly.
-    # We don't use torchao quantization, so a no-op shim is safe.
     try:
         import torch.utils._pytree as _pytree
         if not hasattr(_pytree, "register_constant"):
@@ -72,6 +73,33 @@ def _run_training():
     except Exception as _e:
         _append_log(f"Warning: could not patch _pytree: {_e}")
 
+    patched_dtypes = []
+    dtype_aliases = {
+        **{f"int{i}": torch.int8 for i in range(1, 8)},
+        **{f"uint{i}": torch.uint8 for i in range(1, 8)},
+    }
+    for name, fallback in dtype_aliases.items():
+        if not hasattr(torch, name):
+            setattr(torch, name, fallback)
+            patched_dtypes.append(name)
+
+    if patched_dtypes:
+        _append_log(
+            "Applied torch dtype shims for torchao import: "
+            + ", ".join(patched_dtypes)
+        )
+
+
+def _run_training():
+    global _training_status
+    _training_status = "running"
+    _append_log("Thread started — patching torch/torchao compat then importing unsloth...")
+
+    try:
+        _patch_torchao_import_compat()
+    except Exception as _e:
+        _append_log(f"Warning: could not patch torchao import compat: {_e}")
+
     try:
         import unsloth  # noqa: F401
         _append_log(f"✅ unsloth ready (v{getattr(unsloth, '__version__', 'unknown')})")