Studio: exclude mlx-lm 0.31.3 (broke gemma4/qwen3_5 QK-norm load on Apple Silicon) (#6803)

* Studio: exclude mlx-lm 0.31.3 (broke gemma4/qwen3_5 QK-norm load)

mlx-lm 0.31.3 regressed the QK-norm archs: its strict load_weights rejects the
q_norm/k_norm tensors with "Received N parameters not in model", so gemma4 and
qwen3_5 checkpoints fail to load. Studio installs the MLX stack unpinned at
latest, which pulls 0.31.3. Verified on a real macos-14 runner: gemma4 fails to
load on 0.31.3 but loads and generates coherently on 0.31.2 and on git-main
(future 0.31.4). See mlx-lm #1242.

Exclude just that release (!=0.31.3) in the installer and the self-heal floor so
--upgrade still resolves to the newest good build, and treat an already-installed
0.31.3 as unsatisfied so the self-heal replaces it.

* Studio MLX: cover fresh-install path + robust bad-version compare

Address PR review:
- Fresh install.sh (Apple Silicon) runs the base 'uv pip install unsloth' with
  SKIP_STUDIO_BASE=1, skipping the guarded MLX-stack step, so transitive
  resolution could still pull mlx-lm 0.31.3. install.sh already exports
  UV_OVERRIDE -> overrides-darwin-arm64.txt before that install, so exclude
  mlx-lm 0.31.3 there too; this also strengthens the self-heal (same override).
- Match the known-bad version with parsed packaging.Version so 0.31.3 == 0.31.3.0
  (trailing-zero normalization) instead of raw string equality.

* Studio: exclude mlx-lm 0.31.3 on the fresh Apple Silicon install too

The overrides file only applies via UV_OVERRIDE when it exists relative to the
script, which is not true for a curl-piped install, and the guarded MLX step in
install_python_stack.py is skipped there (SKIP_STUDIO_BASE=1). So the base
install could still resolve the transitive mlx-lm to the broken 0.31.3. Append
mlx-lm!=0.31.3 to the base install on Apple Silicon (empty elsewhere), so the
fresh path pins away from 0.31.3 without waiting for the runtime self-heal.

* Studio: exclude mlx-lm 0.31.3 on the migrated install; keep the >=0.22.0 floor

The with-deps migrated install did not append ${_MLX_LM_EXCLUDE_ARG:-}, so a
curl-piped Apple Silicon migration (no repo overrides file, UV_OVERRIDE unset)
could resolve mlx-lm 0.31.3 transitively. Append the exclusion there, matching
the fresh install path. The no-torch migration is left alone since --no-deps
never resolves mlx-lm (same as the fresh no-torch path).

Also restore the >=0.22.0 floor in overrides-darwin-arm64.txt: a uv override
replaces the transitive constraint, so a bare !=0.31.3 could let the resolver
drop below the supported minimum that mlx_repair.py enforces at runtime.

* Triage huggingface_hub 1.22.0 / fastapi / multiprocess scanner false positives

The scan-packages gate red-failed on all three shards after transitive deps
bumped. Every new CRITICAL is a benign false positive, verified against upstream:

- huggingface_hub 1.22.0 added _sandbox.py for the remote HF sandbox feature.
  Its job-startup bootstrap string (fetch sbx-server into the container /tmp and
  exec it) and the SandboxPool host-reservation loop trip the staged-dropper and
  C2-loop heuristics; that script runs inside a remote HF container, not on the
  user machine. The bump also re-hashed the already-reviewed benign polling loops
  in hf_api.py and utils/_http.py. The PyPI artifact is byte-identical to the
  official v1.22.0 tag.
- fastapi 0.139.0 routing.py re-hashed the websocket keepalive while-True loop;
  byte-identical to upstream 0.139.0.
- multiprocess 0.70.19 forkserver.py and tests/__init__.py re-hashed the AF_UNIX
  fork-server IPC and fd-inheritance tests; genuine uqfoundation release, local
  IPC not network.

Added 7 reviewed allowlist entries (no blind regenerate). All three shards
(hf-stack, studio, extras) exit 0 locally.

* Tighten mlx-lm 0.31.3 exclusion comments

* Trim mlx-lm 0.31.3 exclusion comments

studio/backend/requirements/single-env/overrides-darwin-arm64.txt

@@ -10,3 +10,9 @@ transformers>=4.57.6
 # anyio that also ImportErrors on TaskHandle and 500s the server. An override
 # wins the fight, so force one consistent <4.14 here too.
 anyio<4.14.0
+
+# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5): strict load_weights
+# rejects q_norm/k_norm, so those checkpoints fail to load. mlx-lm #1242.
+# The override also drops it from transitive resolution; keep the >=0.22.0 floor
+# (mirrors mlx_repair.py _MLX_MIN_VERSIONS) or the resolver could go below it.
+mlx-lm>=0.22.0,!=0.31.3

studio/backend/tests/test_mlx_repair.py

@@ -271,6 +271,29 @@ def test_stack_available_requires_runtime_imports_and_versions(monkeypatch):
     assert imported == list(mr._MLX_RUNTIME_IMPORTS)
 
 
+def test_mlx_packages_exclude_known_bad_mlx_lm():
+    # mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5); the install spec
+    # must exclude it so the resolver picks 0.31.2 or >=0.31.4. See mlx-lm #1242.
+    (mlx_lm_spec,) = [p for p in mr.MLX_PACKAGES if p.startswith("mlx-lm")]
+    assert mlx_lm_spec == "mlx-lm>=0.22.0,!=0.31.3"
+
+
+@pytest.mark.parametrize("bad_form", ["0.31.3", "0.31.3.0"])
+def test_known_bad_installed_mlx_lm_triggers_repair(monkeypatch, bad_form):
+    # An installed 0.31.3 counts as unsatisfied so the self-heal replaces it;
+    # parsed-Version compare also catches the trailing-zero form 0.31.3.0.
+    import importlib.metadata as metadata
+
+    def _version(name):
+        return bad_form if name == "mlx-lm" else mr._MLX_MIN_VERSIONS[name]
+
+    monkeypatch.setattr(metadata, "version", _version)
+    monkeypatch.setattr(
+        mr.importlib, "import_module", lambda _n: pytest.fail("versions must gate imports")
+    )
+    assert mr.mlx_stack_available() is False
+
+
 def test_no_op_off_apple_silicon(monkeypatch):
     monkeypatch.setattr(mr, "is_apple_silicon", lambda: False)
     called = {"n": 0}

studio/backend/utils/mlx_repair.py

@@ -45,9 +45,21 @@ DISABLE_ENV_VAR = "UNSLOTH_DISABLE_MLX_AUTOREPAIR"
 # deps). mlx-vlm especially must be >=0.4.4: an older one still imports but
 # breaks VLM Train/Export, so installing it would wrongly clear chat-only.
 _MLX_MIN_VERSIONS = {"mlx": "0.22.0", "mlx-lm": "0.22.0", "mlx-vlm": "0.4.4"}
+# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5): strict load_weights
+# rejects q_norm/k_norm, so a self-heal must not pull it. mlx-lm #1242.
+_MLX_BAD_VERSIONS = {"mlx-lm": ("0.31.3",)}
 _MLX_PACKAGE_NAMES = tuple(_MLX_MIN_VERSIONS)
 _MLX_RUNTIME_IMPORTS = ("mlx.core", "mlx_lm", "mlx_lm.sample_utils", "mlx_vlm")
-MLX_PACKAGES = tuple(f"{name}>={version}" for name, version in _MLX_MIN_VERSIONS.items())
+
+
+def _mlx_spec(name: str, version: str) -> str:
+    spec = f"{name}>={version}"
+    for bad in _MLX_BAD_VERSIONS.get(name, ()):
+        spec += f",!={bad}"
+    return spec
+
+
+MLX_PACKAGES = tuple(_mlx_spec(name, version) for name, version in _MLX_MIN_VERSIONS.items())
 _MLX_REINSTALL_ARGS = tuple(
     arg for name in _MLX_PACKAGE_NAMES for arg in ("--reinstall-package", name)
 )
@@ -140,7 +152,12 @@ def _mlx_versions_satisfy_minimums() -> bool:
         return False
     for name, minimum in _MLX_MIN_VERSIONS.items():
         try:
-            if Version(_dist_version(name)) < Version(minimum):
+            installed = Version(_dist_version(name))
+            if installed < Version(minimum):
+                return False
+            # A known-broken build counts as unsatisfied so the self-heal
+            # reinstalls a good one; Version compare matches 0.31.3(.0/+local).
+            if any(installed == Version(bad) for bad in _MLX_BAD_VERSIONS.get(name, ())):
                 return False
         except PackageNotFoundError:
             return False

studio/install_python_stack.py

@@ -1532,6 +1532,10 @@ LOCAL_DD_UNSTRUCTURED_PLUGIN = (
 )
 LOCAL_DD_GITHUB_PLUGIN = SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed"
 
+# mlx-lm 0.31.3 broke gemma4 / qwen3_5 loading (strict load_weights rejects the
+# QK-norm q_norm/k_norm tensors); exclude just that release. See mlx-lm #1242.
+MLX_LM_BAD_VERSION_EXCLUSION = "!=0.31.3"
+
 # Apple Silicon: override mlx-vlm/mlx-lm's transformers pin (see overrides).
 # _uv_safe_path: uv truncates UV_OVERRIDE at the first space too (issue #6503).
 _MLX_OVERRIDES = SINGLE_ENV / "overrides-darwin-arm64.txt"
@@ -2092,6 +2096,8 @@ def install_python_stack() -> int:
 
     # macOS arm64: install MLX stack at latest (UV_OVERRIDE relaxes the
     # mlx-vlm / mlx-lm transformers pin -- set at module load).
+    # Exclude mlx-lm 0.31.3 (see MLX_LM_BAD_VERSION_EXCLUSION); it broke
+    # gemma4 / qwen3_5 QK-norm loading. mlx-lm #1242.
     if IS_MAC_ARM and not skip_base:
         _progress("MLX stack (Apple Silicon)")
         pip_install(
@@ -2100,7 +2106,7 @@ def install_python_stack() -> int:
             "--upgrade",
             "mlx",
             "mlx-metal",
-            "mlx-lm",
+            f"mlx-lm{MLX_LM_BAD_VERSION_EXCLUSION}",
             "mlx-vlm",
         )