mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-25 08:42:25 +00:00
eb92c60e0b
* Studio: honour the GPU-layer split for DiffusionGemma loads load_model accepted gpu_memory_mode / gpu_layers but the diffusion branch returned early without forwarding them, then overwrote the recorded values with auto / -1. The runner pinned every layer to the GPU, so a GGUF larger than VRAM OOMed in cudaMalloc even with GPU layers set to 0, and /status reported auto regardless of what the user picked. Forward the split to the shim as --ngl in manual mode and record what was actually applied. Auto mode is unchanged. * Address Codex review: dedup, shim compat, and zero-layer residency - Both already-loaded guards skipped the GPU-memory comparison for diffusion, so changing a loaded model from 10 layers to 8 deduped to 'already satisfied' and never restarted the shim. They now compare the EFFECTIVE --ngl via a shared _diffusion_manual_ngl helper, which also avoids a reload loop: the UI keeps a manual preference standing across a diffusion load, so comparing raw modes would reload forever. - --ngl is now gated on _shim_supports_ngl. The unsloth_zoo floor still allows a shim without the flag, where argparse would exit before health and fail the load outright; that case logs an upgrade hint and drops the flag instead. - gpu_layers=0 now CUDA-masks the child like a genuine CPU host. Previously it set _gpu_offload_active=False while still passing a real --gpu token, and the training VRAM coordinator trusts that flag to skip unloading, so a zero-layer runner could hold VRAM into a training run. * Studio: keep the diffusion GPU-layer split honest across every backend Follow-up to the layer-split plumbing in this PR. Four behaviours the split introduced did not hold outside the CUDA happy path. 1. `_gpu_offload_active` was derived from `cpu_only`, which is `_effective_gpu_count() == 0`, which is `torch.cuda` and nothing else. That reads 0 on a Metal, Vulkan, Windows-HIP or Intel XPU host, all of which ship the visual server and will use the GPU, and none of which an empty `CUDA_VISIBLE_DEVICES` can mask (Vulkan selects with GGML_VK_VISIBLE_DEVICES; Metal has no CUDA runtime at all). The training VRAM coordinator treats `False` as "confirmed no VRAM" and skips the unload entirely, so a diffusion runner could survive into a training run holding memory. `False` is now claimed only for an explicit zero-layer split, which is the one case that keeps weights off every backend. 2. `_diffusion_gpu_arg` returned `sorted(gpu_ids)[0]` before it consulted `cpu_only`, so `gpu_layers=0` with a GPU picked still handed the child a real device while the residency flag said otherwise. A new `force_cpu` outranks the picker for a zero-layer request. `cpu_only` deliberately stays below it: it only means no device is visible to torch, and an explicit pick is still the better answer on a host whose GPU torch cannot see. 3. When the installed shim has no `--ngl` the launcher drops the split and records the default, but both dedup guards compared the raw request against that recorded default, so the same manual request mismatched forever and restarted the runner on every /load. `diffusion_requested_ngl` records what was asked for while `gpu_layers` keeps reporting what is running. 4. The active-training preflight never received `gpu_layers`, so a zero-layer diffusion load, which places no layers on any device, was still sized as a full single-GPU occupant and could be refused with 409. The guard now takes the layer count and returns early for that case, and its device probe mirrors the loader's override instead of checking only the torch GPU count. `ValidateModelRequest` gains `gpu_layers` so the preflight sizes the load it is gating, defaulting to -1 so a client that omits it behaves exactly as before. Also parse the shim to detect `--ngl` rather than grepping for the literal `"--ngl"`. The substring test reported support for a shim that merely mentions the flag in a comment, which sends an argument argparse then rejects, and missed a shim that quotes it differently, which silently drops the user's split. The extension check is case-insensitive so a Windows `UNSLOTH_DG_SHIM` pointing at SHIM.PY inspects the file that is actually about to be spawned. Auto mode and the Auto slider are untouched: the child command line and the recorded state are byte-identical to main on every OS and GPU backend. * Tests: make the diffusion GPU-layer contract behavioural Most of these assertions matched source text, so they passed on code that does not work and failed on code that does. `test_zero_layers_masks_the_child_devices` asserted the string `self._gpu_offload_active = not cpu_only` was present, which is why the residency bug above went undetected, and which broke the moment that line was corrected. Both dedup tests were satisfied by the identifier appearing anywhere in the function, so for the route the import line alone passed them and the comparison could have been deleted. The "no longer hardcodes auto" test was two `not in` checks that also pass if the assignments are deleted outright. They now drive the real functions: the device token for a zero-layer request with and without an explicit pick, nine dedup transitions including the unsupported-shim case, and the probe against declarations, comments, docstrings, alternative quoting and an uppercase extension. The training-guard test checks the signature and both call sites. The fixture also left `studio/backend` on `sys.path` for the rest of the session, shadowing generic top-level names (utils, state, models, hub, auth, storage) for every later test in the same run; sibling files avoid this deliberately. It is popped in a `finally` now. 16 tests to 35. * Use the AGPL-3.0 header on the new studio test, matching the tree * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Re-trigger CI The merge push and the pre-commit.ci formatting push landed two minutes apart, so the concurrency groups cancelled all 46 checks on this head before any of them reported. * Address the second Codex round: guard/launcher agreement and hydration The training guard and validate path now reason about what the launcher will actually run, not what was asked: - The zero-layer bypass and the new positive-split scaling are gated on diffusion_split_supported(), a mirror of the launch gate. A split the shim cannot accept is guarded as the default GPU-resident configuration, so a manual/0 request against an old unsloth_zoo no longer skips the VRAM check while the child takes a whole GPU. - A confirmed-diffusion positive split is sized as its GPU share (ngl/n_layers of the estimate) instead of the whole file, so gpu_layers=20 during training no longer 409s a load that fits. Unknown layer count keeps the full, conservative estimate. - _shim_supports_ngl keys on the argv shape instead of a .py suffix: an extensionless or .pyw UNSLOTH_DG_SHIM override is the file about to be spawned and is now the file probed. - Both frontend validate call sites now send gpu_layers alongside gpu_memory_mode, so validate reaches the same verdict as the /load that follows it. - Status hydration adopts gpuMemoryMode for a diffusion runner that reports manual (a split was actually applied); a diffusion auto response still leaves the standing preference alone, so the old no-loop behaviour holds. * Address the third Codex round: probe precision, dedup after upgrade, more call sites - _shim_supports_ngl no longer consults a sibling shim.py next to a file override: the override itself is what argparse will run, and a capable sibling vouching for an incapable custom-named shim re-created the exact startup failure the probe exists to prevent. The extra_pythonpath parameter is gone; it could only ever name that sibling. - A dropped split now un-sticks after a mid-session unsloth_zoo upgrade: both dedup guards re-probe shim support when the recorded state shows the ask was not applied, so an identical Apply finally reloads with --ngl. The probe only runs in that narrow state, keeping the no-loop behaviour against a shim that stays old. - chat-adapter's background auto-load preflight now sends gpu_layers, the third validate call site: a remembered manual split (0 in particular) was still being sized as a full-GGUF occupant during training. - Compare panes no longer force a diffusion pane to auto/-1: the runner honours the layer split now, so panes pass gpuMemoryMode/gpuLayers through like every other GGUF. MoE offload and tensor parallel stay forced off; the runner has no equivalent knobs. * Keep the standing diffusion layer count across a dropped-split hydration When an older shim drops a manual split, the backend reports auto/-1 (what ran) while keeping the ask in diffusion_requested_ngl. Hydration preserved the standing gpuMemoryMode for that response but still reset the editable gpuLayers to Auto, so the next Apply sent manual/-1 and the ask could never be applied, even after the unsloth_zoo upgrade that adds --ngl. The non-manual reset now skips gpuLayers for diffusion, mirroring the standing mode; loaded baselines and the MoE/split knobs still reset. * Do not leak another model's GPU split into a diffusion compare pane main pinned a diffusion compare pane to auto/-1: const effectiveGpuMemoryMode = resolvedIsDiffusion ? "auto" : (ownConfig.gpuMemoryMode ?? compareLoadKnobs.gpuMemoryMode); This PR removed the pin so a pane's own saved split can be honoured, but the replacement is an unguarded ?? chain and DEFAULT_PER_MODEL_CONFIG defines neither gpuMemoryMode nor gpuLayers. An unconfigured diffusion pane therefore falls through to compareLoadKnobs, which is a snapshot of the live store, i.e. whatever chat GGUF was loaded at Send, whose layer count is bounded by that model. Repro in one session: set Manual + 12 layers on a normal chat GGUF, open compare, add a never-configured DiffusionGemma, Send. The runner launches with --ngl 12. A standing 0 is worse: force_cpu masks the pane's devices and it runs entirely on CPU without being asked. No diffusion UI can show or clear the inherited value, because the mode row and the layer slider are hidden for diffusion. The pane's OWN split is still sent, so the feature this PR adds is intact; only the inherited snapshot is refused. nCpuMoe and tensorParallel two lines below already guard on resolvedIsDiffusion the same way. The rule moves to a dependency-free module so it can be unit-tested: the frontend runner is node --experimental-strip-types with no @/ alias resolution, so chat-runtime-store.ts is not importable from a test. GPU_LAYERS_AUTO moves with it and is re-exported, so every existing import path keeps working. * Only bypass the training guard for a confirmed diffusion zero-layer split The zero-layer bypass used `diffusion_kind is not False`, which also exempts an UNCLASSIFIED GGUF. _classify_diffusion_gguf returns None whenever the header is unreadable and the name lacks the family, so an uncached HF-repo GGUF loaded at manual/0 while training is active skipped the VRAM estimate entirely. That is unsafe for an ordinary GGUF: --gpu-layers 0 is not zero VRAM when a device pin, surviving tensor mode, an mmproj or a GPU drafter keeps it resident, which is exactly what LlamaCppBackend._zero_offload_keeps_gpu_visible exists to report. The scaling block sixty lines below already refuses an unknown classification for the same reason, and says so in its comment, so the two were contradicting each other. A DiffusionGemma the user actually downloaded classifies True either by name or by header, so nothing this PR adds is lost. Adds two behavioural cases: a confirmed diffusion split at ngl 0 still bypasses, an unclassified GGUF at ngl 0 still raises 409. The second fails against the old condition. * Keep an unclassified GGUF on diffusion-safe compare placement _classify_diffusion_gguf is a tri-state, but /validate collapsed None into is_diffusion = false. An undownloaded GGUF has no header to read, so one whose repo/file name does not carry the DiffusionGemma family came back looking exactly like a confirmed ordinary GGUF, and a compare pane with no own config inherited the active model's manual split. /load then downloads the file, reads a diffusion header and applies that split anyway: an inherited 0 CPU-masks the runner, and another partial count repartitions it or OOMs. /validate now reports the inconclusive case as diffusion_unknown, and shouldPinDiffusionPlacement pins an unclassified GGUF to Auto exactly like a confirmed diffusion one. is_diffusion keeps its old meaning, so the other three preflight consumers are untouched, and the new field is additive: absent on an older backend it reads as "classified", which is today's behaviour. A readable non-diffusion header still inherits the snapshot, so the common path is unchanged. * Point the compare-placement contract at the module that now owns the rule Extracting the pane's mode/layer fallback into lib/gpu-placement.ts left test_compare_load_uses_each_models_gpu_config asserting two source substrings that no longer exist in shared-composer.tsx, so Repo tests (CPU) would have failed. Assert the delegation here and the fallback at its new home; the rule itself is covered behaviourally by studio/frontend/tests/gpu-placement.test.ts. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten the diffusion GPU-layer comments Condense the comments and docstrings added by this PR: drop the restatements, collapse the multi-line explanations that still read clearly on one line, and keep the reasoning that is not obvious from the code (the force_cpu precedence, why the shim is parsed rather than grepped, and why the dedup guards compare the requested split instead of the applied one). Comments only, no behaviour change. * Keep the unknown diffusion verdict through a config-picker selection The compare pane re-probes only when sel.isDiffusion is undefined, but the model-config page collapsed an inconclusive probe to a definite false before handing it to onRun: classifiedIsDiffusion = isDiffusion ? true : stagedDims?.isDiffusion For an uncached GGUF whose name carries no family that is false, which flows through the selection into the pane, skips the re-probe, and leaves diffusionUnknown false. The pane then inherits the other model's manual layer count, which /load applies once the downloaded header turns out to be diffusion. Same leak as the last fix, through a different entry point. resolveStagedDiffusionClassification keeps the answer tri-state, so an inconclusive probe stays undefined and the pane's existing preflight re-probes and learns diffusion_unknown. A conclusive answer is unchanged in both directions, and the page's own resolvedIsDiffusion still tests === true, so no GPU control changes visibility. * Report the diffusion split a shim dropped, so a refresh can recover it A shim without --ngl cannot apply a manual split: the runner reports Auto while the backend keeps the ask in _diffusion_requested_ngl. That value never reached the client, so the recovery only worked while the in-memory store survived. After a browser refresh the store starts at Auto and /status reports Auto, the ask is gone, and the next Apply sends manual/-1 -- which the retry that exists to apply the split after an unsloth_zoo upgrade can never turn back into the original count. LoadResponse and InferenceStatusResponse now carry diffusion_requested_ngl, and recoverDroppedDiffusionSplit restores it during hydration, re-asserting manual mode so the next Apply can re-send the count. Zero is preserved as the real CPU-only ask it is. Both fields default to None and the frontend reads them optionally, so a backend without them behaves exactly as before. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Stop budgeting an unclassified zero-layer GGUF as CPU-only The confirmed-diffusion early return was hardened to `is True`, but the force_cpu on the _diffusion_gpu_arg call sixty lines below still keyed off `diffusion_ngl == 0` alone. An unclassified GGUF therefore skipped the early return, as intended, and then handed the guard an empty device token anyway. can_load_chat_during_training reads an empty single_device_gpu as a CPU-only runner and returns True unconditionally, so the load was allowed during training on an assumption that only holds for a real diffusion model. The comment above the early return already gives the reason it does not: a device pin, tensor mode, mmproj or a GPU drafter keeps an ordinary GGUF resident at --gpu-layers 0. Drop force_cpu at this call site rather than gating it. Reaching it with ngl 0 now implies an unclassified GGUF, since the confirmed case returns above, so the flag could only ever be False. The picker then chooses a device and the load stays conservatively sized. The existing unclassified-zero test could not catch this: it stubs can_load, so the empty-token branch never ran, and it asserted only that the guard was reached. The new test asserts the token it was reached with. * Tighten the comments added since the last comment pass --------- Co-authored-by: Daniel Han <unslothai@gmail.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
114 lines
4.5 KiB
Python
114 lines
4.5 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""/api/inference/validate must report an INCONCLUSIVE diffusion check as such.
|
|
|
|
`_classify_diffusion_gguf` is a tri-state: True (diffusion), False (header read,
|
|
ordinary), None (nothing to read, and no family in the name).
|
|
|
|
The response used to collapse None into `is_diffusion = False`, so a caller could not
|
|
tell "ordinary GGUF" from "unknown". The staged-metadata preflight picks a GPU-layer
|
|
split from that answer, and /load may then apply it to a diffusion runner: an inherited
|
|
0 CPU-masks it, another count repartitions or OOMs it.
|
|
"""
|
|
|
|
import asyncio
|
|
import importlib.util
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from models.inference import ValidateModelRequest
|
|
|
|
|
|
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
|
|
|
|
|
|
def _load_route_module(name: str):
|
|
spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / "routes/inference.py")
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
async def _noop_gpu_ids(_config, gpu_ids, **_kwargs):
|
|
return gpu_ids, False
|
|
|
|
|
|
class TestValidateReportsDiffusionUnknown(unittest.TestCase):
|
|
def _validate(
|
|
self,
|
|
route,
|
|
*,
|
|
diffusion_kind,
|
|
is_gguf = True,
|
|
):
|
|
# Mirrors the real staged-metadata preflight; also skips the training guard.
|
|
request = ValidateModelRequest(
|
|
model_path = "someone/repacked-gguf",
|
|
include_context_length = True,
|
|
)
|
|
config = SimpleNamespace(
|
|
identifier = "someone/repacked-gguf",
|
|
display_name = "repacked-gguf",
|
|
is_gguf = is_gguf,
|
|
is_lora = False,
|
|
is_vision = False,
|
|
gguf_file = None,
|
|
)
|
|
with (
|
|
patch.object(
|
|
route,
|
|
"_resolve_model_identifier_for_request",
|
|
return_value = ("someone/repacked-gguf", "someone/repacked-gguf", False),
|
|
),
|
|
patch.object(route.ModelConfig, "from_identifier", return_value = config),
|
|
patch.object(route, "_resolve_inherited_extra_args", return_value = None),
|
|
patch.object(route, "_classify_diffusion_gguf", return_value = diffusion_kind),
|
|
patch.object(route, "_resolve_gguf_gpu_ids_for_request", new = _noop_gpu_ids),
|
|
patch.object(route, "_effective_load_in_4bit", return_value = True),
|
|
):
|
|
return asyncio.run(route.validate_model(request, current_subject = "test-user"))
|
|
|
|
def test_unclassifiable_gguf_is_reported_unknown_not_ordinary(self):
|
|
"""The bug: None must not look identical to a confirmed ordinary GGUF."""
|
|
route = _load_route_module("inf_route_diffusion_unknown_1")
|
|
resp = self._validate(route, diffusion_kind = None)
|
|
self.assertFalse(resp.is_diffusion)
|
|
self.assertTrue(
|
|
resp.diffusion_unknown,
|
|
"an unreadable/undownloaded GGUF with no family in its name is UNKNOWN; "
|
|
"reporting it as a plain non-diffusion GGUF lets a caller inherit a "
|
|
"GPU-layer split that /load will apply to a diffusion runner",
|
|
)
|
|
|
|
def test_confirmed_ordinary_gguf_is_not_unknown(self):
|
|
route = _load_route_module("inf_route_diffusion_unknown_2")
|
|
resp = self._validate(route, diffusion_kind = False)
|
|
self.assertFalse(resp.is_diffusion)
|
|
self.assertFalse(resp.diffusion_unknown)
|
|
|
|
def test_confirmed_diffusion_gguf_is_not_unknown(self):
|
|
route = _load_route_module("inf_route_diffusion_unknown_3")
|
|
resp = self._validate(route, diffusion_kind = True)
|
|
self.assertTrue(resp.is_diffusion)
|
|
self.assertFalse(resp.diffusion_unknown)
|
|
|
|
def test_non_gguf_is_never_unknown(self):
|
|
"""A transformers model is definitively not a diffusion GGUF."""
|
|
route = _load_route_module("inf_route_diffusion_unknown_4")
|
|
resp = self._validate(route, diffusion_kind = False, is_gguf = False)
|
|
self.assertFalse(resp.is_diffusion)
|
|
self.assertFalse(resp.diffusion_unknown)
|
|
|
|
def test_flag_defaults_off_so_an_old_client_reads_the_same_response(self):
|
|
"""Additive field: absent/False keeps the pre-#7575 meaning of is_diffusion."""
|
|
from models.inference import ValidateModelResponse
|
|
|
|
resp = ValidateModelResponse(valid = True, message = "ok")
|
|
self.assertFalse(resp.diffusion_unknown)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|