Fix ZeroGPU v14p native Transformers runtime

figment/zerogpu_runtime.py

@@ -55,6 +55,7 @@ def _runtime_for(model_repo: str, model_subfolder: str, model_id: str) -> "_Zero
 class _ZeroGpuRuntime:
     def __init__(self, *, model_repo: str, model_subfolder: str, model_id: str) -> None:
         import torch
+        from transformers import AutoConfig
         from transformers import AutoModelForCausalLM
         from transformers import AutoTokenizer
 
@@ -66,14 +67,21 @@ class _ZeroGpuRuntime:
         self.tokenizer = AutoTokenizer.from_pretrained(
             model_repo,
             subfolder=model_subfolder,
-            trust_remote_code=True,
+            trust_remote_code=False,
         )
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
+        model_config = AutoConfig.from_pretrained(
+            model_repo,
+            subfolder=model_subfolder,
+            trust_remote_code=False,
+        )
+        model_config.use_mamba_kernels = False
         self.model = AutoModelForCausalLM.from_pretrained(
             model_repo,
             subfolder=model_subfolder,
-            trust_remote_code=True,
+            config=model_config,
+            trust_remote_code=False,
             torch_dtype=torch.bfloat16,
         )
         self.model.to("cuda")

requirements.txt

@@ -1,13 +1,13 @@
 fastapi>=0.115,<1
 gradio>=6.0,<7
-huggingface_hub>=1.0,<2
+huggingface_hub>=1.5,<2
 httpx>=0.27,<1
 pydantic>=2.7,<3
 python-dotenv>=1.0,<2
 PyYAML>=6.0,<7
 spaces>=0.40,<1
 torch>=2.8,<2.12
-transformers>=4.52,<5
+transformers>=5.10,<6
 accelerate>=1.8,<2
 einops>=0.8,<1
 protobuf>=5,<7
@@ -15,5 +15,3 @@ safetensors>=0.5,<1
 sentencepiece>=0.2,<1
 setuptools>=70
 wheel>=0.45
-causal-conv1d>=1.6,<2
-mamba-ssm>=2.2,<3

tests/test_zerogpu_runtime.py

@@ -0,0 +1,64 @@
+import sys
+from types import ModuleType
+from types import SimpleNamespace
+from typing import Any
+
+from figment.zerogpu_runtime import _ZeroGpuRuntime
+
+
+def test_zerogpu_runtime_uses_native_transformers_with_mamba_kernels_disabled(monkeypatch: Any) -> None:
+    calls: dict[str, Any] = {}
+    fake_config = SimpleNamespace(use_mamba_kernels=True, use_cache=False)
+
+    fake_torch = ModuleType("torch")
+    fake_torch.bfloat16 = "bf16"
+
+    class FakeAutoConfig:
+        @staticmethod
+        def from_pretrained(*args: Any, **kwargs: Any) -> Any:
+            calls["config"] = {"args": args, "kwargs": kwargs}
+            return fake_config
+
+    class FakeTokenizer:
+        pad_token = "<pad>"
+        eos_token = "</s>"
+
+    class FakeAutoTokenizer:
+        @staticmethod
+        def from_pretrained(*args: Any, **kwargs: Any) -> FakeTokenizer:
+            calls["tokenizer"] = {"args": args, "kwargs": kwargs}
+            return FakeTokenizer()
+
+    class FakeModel:
+        def __init__(self) -> None:
+            self.config = SimpleNamespace(use_cache=False)
+
+        def to(self, device: str) -> None:
+            calls["device"] = device
+
+        def eval(self) -> None:
+            calls["eval"] = True
+
+    class FakeAutoModelForCausalLM:
+        @staticmethod
+        def from_pretrained(*args: Any, **kwargs: Any) -> FakeModel:
+            calls["model"] = {"args": args, "kwargs": kwargs}
+            return FakeModel()
+
+    fake_transformers = ModuleType("transformers")
+    fake_transformers.AutoConfig = FakeAutoConfig
+    fake_transformers.AutoModelForCausalLM = FakeAutoModelForCausalLM
+    fake_transformers.AutoTokenizer = FakeAutoTokenizer
+    monkeypatch.setitem(sys.modules, "torch", fake_torch)
+    monkeypatch.setitem(sys.modules, "transformers", fake_transformers)
+
+    _ZeroGpuRuntime(model_repo="repo", model_subfolder="subfolder", model_id="model")
+
+    assert calls["tokenizer"]["kwargs"]["trust_remote_code"] is False
+    assert calls["config"]["kwargs"]["trust_remote_code"] is False
+    assert fake_config.use_mamba_kernels is False
+    assert calls["model"]["kwargs"]["config"] is fake_config
+    assert calls["model"]["kwargs"]["trust_remote_code"] is False
+    assert calls["model"]["kwargs"]["torch_dtype"] == "bf16"
+    assert calls["device"] == "cuda"
+    assert calls["eval"] is True