utils

training/chat_utils.py

@@ -0,0 +1,50 @@
+"""Chat rendering helpers for text-only Qwen/Qwen-VL control prompts."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+
+def render_no_think_chat(
+    tokenizer: Any,
+    messages: List[Dict[str, str]],
+    *,
+    add_generation_prompt: bool,
+) -> str:
+    """Render a chat prompt with Qwen thinking disabled when supported.
+
+    Qwen3-family templates expose ``enable_thinking`` as a Jinja variable.
+    Older templates ignore that keyword, so we fall back cleanly rather than
+    failing training for non-Qwen or older tokenizer builds.
+    """
+    kwargs = {
+        "tokenize": False,
+        "add_generation_prompt": add_generation_prompt,
+        "enable_thinking": False,
+    }
+    try:
+        return tokenizer.apply_chat_template(messages, **kwargs)
+    except TypeError as exc:
+        if "enable_thinking" not in str(exc):
+            raise
+        kwargs.pop("enable_thinking")
+        return tokenizer.apply_chat_template(messages, **kwargs)
+
+
+def tokenize_text_only(tokenizer: Any, input_text: str, device: Any):
+    """Tokenize a rendered text prompt without invoking VL image loading.
+
+    Some Qwen-VL processors route the first positional argument to ``images``.
+    Passing the transcript through the explicit ``text=`` keyword keeps the
+    prompt on the text path and avoids PIL trying to parse chat text as an image.
+    """
+    try:
+        inputs = tokenizer(text=input_text, return_tensors="pt")
+    except ValueError as exc:
+        if "Incorrect image source" not in str(exc):
+            raise
+        inputs = tokenizer(text=[input_text], images=None, return_tensors="pt")
+    except TypeError:
+        inputs = tokenizer(input_text, return_tensors="pt")
+
+    return inputs.to(device)

training/launch_smoke.py

@@ -46,9 +46,8 @@ def build_job_command() -> str:
         "cd /workspace/AntiAtropos\n"
         "\n"
         "echo '[bootstrap] Installing dependencies...'\n"
-        "pip install --break-system-packages \\\n"
-        "  torch==2.10.0+cu124 torchvision==0.20.0+cu124 torchaudio==2.10.0+cu124 \\\n"
-        "  --index-url https://download.pytorch.org/whl/cu124 -q\n"
+        "pip install --break-system-packages --no-deps torchvision -q\n"
+        "pip install --break-system-packages flash-attn --no-build-isolation -q\n"
         "pip install --break-system-packages -r training/requirements.txt -q\n"
         "\n"
         "echo '[bootstrap] Launching training...'\n"

training/requirements.txt

@@ -2,11 +2,15 @@
 # Install: pip install -r training/requirements.txt
 
 # ---- Core ML ----
-# NOTE: torch is provided by the Docker base image (pytorch/pytorch:2.11.0).
+# NOTE: torch is provided by the Docker base image (pytorch/pytorch:2.10.0).
 # Do NOT list it here — pip's resolver may downgrade it, breaking torchao.
 # torch>=2.5.0  (pre-installed in base image)
-# torchao is also pre-installed in the 2.11.0 base image and compatible.
-torchvision  # required by unsloth_zoo.vision_utils
+# torchao is also pre-installed in the base image and compatible.
+# torchvision is required by unsloth_zoo.vision_utils but MUST be installed
+# separately via the PyTorch index URL (not via pip resolver) to avoid
+# torch version downgrades. Install before running this file:
+#   pip install torch==X.Y.Z+cuNNN torchvision==A.B.C+cuNNN torchaudio==D.E.F+cuNNN \
+#     --index-url https://download.pytorch.org/whl/cuNNN
 transformers>=4.45.0
 accelerate>=0.34.0
 bitsandbytes>=0.43.0

training/run_smoke_uv.py

@@ -37,6 +37,10 @@ subprocess.run(
 os.chdir(str(WORKSPACE / "AntiAtropos"))
 
 print("[bootstrap] Installing full dependencies...")
+subprocess.run(
+    ["uv", "pip", "install", "--no-deps", "torchvision", "-q"],
+    check=True,
+)
 subprocess.run(
     ["uv", "pip", "install", "-r", "training/requirements.txt", "-q"],
     check=True,