Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
danielhanchen merged 9 commits into from
Jul 7, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
c36248e
Studio: exclude mlx-lm 0.31.3 (broke gemma4/qwen3_5 QK-norm load)
danielhanchen Jul 2, 2026
aaee1b2
Studio MLX: cover fresh-install path + robust bad-version compare
danielhanchen Jul 2, 2026
ce773ed
Studio: exclude mlx-lm 0.31.3 on the fresh Apple Silicon install too
danielhanchen Jul 2, 2026
82e2ee2
Studio: exclude mlx-lm 0.31.3 on the migrated install; keep the >=0.2…
danielhanchen Jul 3, 2026
4847bbd
Merge remote-tracking branch 'origin/main' into pr6803-main-update
danielhanchen Jul 3, 2026
7888043
Triage huggingface_hub 1.22.0 / fastapi / multiprocess scanner false …
danielhanchen Jul 5, 2026
81ef81e
Tighten mlx-lm 0.31.3 exclusion comments
danielhanchen Jul 6, 2026
4322dc3
Trim mlx-lm 0.31.3 exclusion comments
danielhanchen Jul 6, 2026
4a23977
Merge main into studio-mlx-arch-version-gate
danielhanchen Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions install.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1442,8 +1442,14 @@ if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then
SKIP_TORCH=true
fi

# Apple Silicon: exclude broken mlx-lm 0.31.3 (QK-norm load regression for
# gemma4 / qwen3_5; mlx-lm #1242). A curl-piped install has no overrides file
# and skips the guarded MLX step (SKIP_STUDIO_BASE=1), so this is the only cover.
_MLX_LM_EXCLUDE_ARG=""

# Apple Silicon: override mlx-vlm / mlx-lm's transformers pin (see overrides file).
if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
_MLX_LM_EXCLUDE_ARG="mlx-lm!=0.31.3"
_OVERRIDES_FILE="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)/studio/backend/requirements/single-env/overrides-darwin-arm64.txt"
if [ -f "$_OVERRIDES_FILE" ]; then
# uv splits UV_OVERRIDE on whitespace, so a repo path with whitespace
Expand Down Expand Up @@ -2679,9 +2685,11 @@ if [ "$_MIGRATED" = true ]; then
run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
fi
else
# Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no
# overrides file, so UV_OVERRIDE is unset and this positional is the only cover.
run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" ${_MLX_LM_EXCLUDE_ARG:-}
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
Expand Down Expand Up @@ -2912,7 +2920,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
else
run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth -- "$PACKAGE_NAME"
--upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-}

@chatgpt-codex-connector chatgpt-codex-connector Bot Jul 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Apply the MLX exclusion on migrations too

This appends mlx-lm!=0.31.3 only on the fresh non-migrated install path. In the legacy migration path above, a healthy $STUDIO_HOME/.venv is moved into $VENV_DIR, _MIGRATED=true, and the installer runs the separate install unsloth (migrated) command without this extra requirement before invoking setup with SKIP_STUDIO_BASE=1; for curl-piped Apple Silicon migrations there is also no repo-relative overrides file to export. That path can still resolve a missing/transitive mlx-lm to the known-bad 0.31.3 and leave QK-norm MLX models broken until runtime autorepair succeeds.

Useful? React with 👍 / 👎.

fi
# AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in
# CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1.
Expand Down
16 changes: 16 additions & 0 deletions scripts/scan_packages_baseline.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1545,6 +1545,22 @@
"severity": "HIGH",
"evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)",
"evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9"
},
{
"package": "multiprocess",
"file": "multiprocess/forkserver.py",
"check": "Reverse shell / bind shell pattern",
"severity": "CRITICAL",
"evidence": "L6: import socket sha256:6c707119169286c9a798e2c8d13a48614e481d8a503950916fd4ffb4c94d3182",
"evidence_hash": "50fec0f0522a8e4e636bf348b752002d7935d8455af31fb78c6f11e2eba19f6d"
},
{
"package": "multiprocess",
"file": "multiprocess/tests/__init__.py",
"check": "Reverse shell / bind shell pattern",
"severity": "CRITICAL",
"evidence": "L3569: os.dup2(conn.fileno(), i) | L3601: \"test needs os.dup2()\") | L3619: os.dup2(fd, newfd) | L20: import socket sha256:c824dc0f409f242420c3fbb324790c53cb3078d2c8b07ee8f2a05694b01c2946",
"evidence_hash": "3878a2b430c175dbc5877a95195bfe52f9588ff73fb74e2261ed5e33087915ad"
}
]
}
6 changes: 6 additions & 0 deletions studio/backend/requirements/single-env/overrides-darwin-arm64.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

@chatgpt-codex-connector chatgpt-codex-connector Bot Jul 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Preserve mlx-vlm's mlx-lm floor when excluding 0.31.3

For Apple Silicon installs using UV_OVERRIDE, this is not an additive constraint: uv pip install --help says overrides are absolute and completely replace package requirements. Current PyPI metadata for mlx-vlm 0.6.3 (https://pypi.org/project/mlx-vlm/0.6.3/) requires mlx-lm>=0.31.3, so replacing that with mlx-lm>=0.22.0,!=0.31.3 allows uv to resolve latest mlx-vlm with mlx-lm 0.31.2 after the bad release is excluded; the self-heal then accepts that upstream-invalid stack because it only checks the same 0.22.0 floor, leaving VLM load/training paths running against a dependency combination mlx-vlm explicitly rejected. Constrain/downgrade mlx-vlm alongside the exclusion, or keep the upstream floor once a good mlx-lm>=0.31.4 exists.

Useful? React with 👍 / 👎.

@danielhanchen danielhanchen Jul 6, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The override deliberately keeps the only enforceable floor (mlx-lm at least 0.22.0) and uv resolves mlx-lm 0.31.2, the last good release before the QK-norm regression and exactly the version this PR targets. mlx-vlm 0.6.3's declared floor (mlx-lm 0.31.3 or newer) is unenforceable here: 0.31.3 is the broken release being excluded and nothing newer exists, so enforcing it yields an empty set and fails the install. There is no reachable path to a too-old or broken mlx-lm, so the enforceable floor is already correctly preserved.

23 changes: 23 additions & 0 deletions studio/backend/tests/test_mlx_repair.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -271,6 +271,29 @@ def _import_module(name):
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}
Expand Down
21 changes: 19 additions & 2 deletions studio/backend/utils/mlx_repair.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,9 +45,21 @@
# 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)
)
Expand Down Expand Up @@ -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
Expand Down
8 changes: 7 additions & 1 deletion studio/install_python_stack.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1532,6 +1532,10 @@ def _torch_step_label(suffix: str) -> str:
)
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"
Expand Down Expand Up @@ -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(
Expand All @@ -2100,7 +2106,7 @@ def install_python_stack() -> int:
"--upgrade",
"mlx",
"mlx-metal",
"mlx-lm",
f"mlx-lm{MLX_LM_BAD_VERSION_EXCLUSION}",

@chatgpt-codex-connector chatgpt-codex-connector Bot Jul 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Cover the fresh install path too

On the normal install.sh Apple Silicon install path, install.sh first runs uv pip install --python "$_VENV_PY" -- "$PACKAGE_NAME" (lines 2899-2900) and then invokes setup.sh with SKIP_STUDIO_BASE=1 (lines 3010 and 3042), so this guarded MLX stack step is skipped entirely. That means the new mlx-lm!=0.31.3 constraint only protects standalone unsloth studio update / non-skip-base runs; fresh installs can still resolve the transitive MLX dependency to the known-bad 0.31.3 and hit the same QK-norm load failure until the best-effort runtime autorepair succeeds. Please add the same exclusion to the install.sh unsloth install path or ensure this MLX repair step runs when skip_base is set on Apple Silicon.

Useful? React with 👍 / 👎.

"mlx-vlm",
)

Expand Down
Loading