vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-09 15:58:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

Studio: allow CPU-only DiffusionGemma by granting the diffusion runner the CPU device (#6979)

Browse source 
* Studio: allow CPU-only DiffusionGemma by granting the diffusion runner the CPU device

* Studio: mark CPU-only DiffusionGemma as non-GPU-resident for training VRAM preflight

* Studio: keep the CPU DiffusionGemma change minimal (revert VRAM-flag tweak; Metal hosts still hold unified memory)

* Studio: keep CPU DiffusionGemma fallback fully CPU-masked so a masked GPU host does not re-expose GPU 0
This commit is contained in:
Daniel Han 2026-07-08 07:26:10 -07:00 • committed by GitHub
parent 3d41e5868d
commit 116ce48c1a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 11 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

12
studio/backend/core/inference/llama_cpp.py
View file

@ -3971,7 +3971,11 @@ class LlamaCppBackend:
# Auto-size (0): the visual server probes the largest context that fits this GPU's VRAM
# (capped at the training context). An explicit in-range n_ctx overrides it.
maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 65536) else 0
gpu = os.environ.get("DG_GPU", "0")
# No visible CUDA GPU: a genuine CPU host, or a GPU host masked with
# CUDA_VISIBLE_DEVICES="" to force CPU serving. Keep the visual-server child
# CPU-masked (empty --gpu) so the shim does not re-expose GPU 0 via its default.
cpu_only = self._effective_gpu_count() == 0
gpu = "" if cpu_only else os.environ.get("DG_GPU", "0")
cmd = list(shim_cmd) + [
"--gguf",
@ -3991,6 +3995,12 @@ class LlamaCppBackend:
# refuses to load unless UNSLOTH_IS_PRESENT is set (normally by `import
# unsloth`). The shim never imports unsloth, so set it here as unsloth does.
env["UNSLOTH_IS_PRESENT"] = "1"
# The shim's `import unsloth_zoo` aborts in get_device_type() ("needs a GPU")
# when no accelerator is visible, even though it only drives the CPU
# visual-server binary and does no torch GPU work. Allow the CPU device so the
# runner starts; the visual server still runs on the CPU llama.cpp build.
if cpu_only:
env.setdefault("UNSLOTH_ALLOW_CPU", "1")
env["DG_VISUAL_BIN"] = visual_bin
env["DG_GPU"] = gpu
# The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH.