Small files from the Hub, shards from the mount (size-checked); drop flash-linear-attention (no GPU at import on ZeroGPU)

app.py

@@ -1,9 +1,10 @@
 """Chat with Whittle-Qwen-3.8-45B-A3B on ZeroGPU.
 
-The model repo is mounted read-only at /models/w45 (no 90 GB download into the Space). The 35 B body runs on a full
+The model repo is mounted read-only at /models/w45 (no 90 GB download into the Space); small files come from the Hub. The 35 B body runs on a full
 RTX PRO 6000 (xlarge); the 10 B n-gram table stays in CPU memory, as with llama.cpp's `-ot per_layer_token_embd=CPU`.
 """
 import os
+import shutil
 import threading
 
 import spaces  # must come before torch/CUDA use on ZeroGPU
@@ -15,8 +16,31 @@ from whittle_load import load_whittle
 
 REPO = "logic65/Whittle-Qwen-3.8-45B-A3B"
 MOUNT = "/models/w45"
-SOURCE = MOUNT if os.path.isfile(os.path.join(MOUNT, "config.json")) else REPO
-
+LOCAL = "/tmp/w45"
+
+
+def prepare_local_dir() -> str:
+    """Small files (config, index, tokenizer, template) come straight from the Hub; the weight shards are read from the
+    read-only mount through symlinks, after checking every shard's size on the mount against the Hub."""
+    from huggingface_hub import HfApi, hf_hub_download
+    files = {f.path: f.size for f in HfApi().list_repo_tree(REPO) if "/" not in f.path and getattr(f, "size", None) is not None}
+    os.makedirs(LOCAL, exist_ok=True)
+    for name, size in sorted(files.items()):
+        dst = os.path.join(LOCAL, name)
+        if name.endswith(".safetensors"):
+            src = os.path.join(MOUNT, name)
+            got = os.path.getsize(src) if os.path.exists(src) else None
+            print(f"[mount] {name}: Hub {size:,} B, mount {got if got is None else f'{got:,} B'}", flush=True)
+            if got != size:
+                raise RuntimeError(f"{name} on the mount does not match the Hub ({got} vs {size} bytes)")
+            if not os.path.lexists(dst):
+                os.symlink(src, dst)
+        elif not name.endswith((".md", ".svg", ".py")) and name != ".gitattributes":
+            shutil.copy(hf_hub_download(REPO, name), dst)
+    return LOCAL
+
+
+SOURCE = prepare_local_dir()
 tokenizer = AutoTokenizer.from_pretrained(SOURCE)
 model = load_whittle(SOURCE, dtype=torch.bfloat16, device_map="cuda").eval()
 

requirements.txt

@@ -3,4 +3,3 @@ transformers>=5.18.0
 accelerate
 safetensors
 huggingface_hub>=1.0
-flash-linear-attention==0.5.2