Studio: allow CPU-only DiffusionGemma by granting the diffusion runner the CPU device - #6979
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates studio/backend/core/inference/llama_cpp.py to set the UNSLOTH_ALLOW_CPU environment variable to '1' when running on a GPU-less host. This prevents the shim's import unsloth_zoo from aborting on CPU-only systems. There are no review comments to address, and I have no additional feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: acd4492b8e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if self._effective_gpu_count() == 0: | ||
| env.setdefault("UNSLOTH_ALLOW_CPU", "1") |
There was a problem hiding this comment.
When this new CPU-only branch is taken, the runner is allowed to start with no visible GPU, but _start_diffusion_server still later sets _gpu_offload_active = True before and after the health check. routes/training_vram.py treats any active GGUF with this field not False as consuming VRAM and unloads/reports it during training preflight, so a CPU-only DiffusionGemma session is incorrectly considered a GPU-resident model. This path should also publish the diffusion server as CPU-only, e.g. by setting _gpu_offload_active to False when no effective GPU is present.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch. Fixed in ba50254: the CPU-only branch now sets _gpu_offload_active = False (via not cpu_only), matching the llama-server path (_classify_gpu_offload), so training_vram.py no longer counts a CPU-only DiffusionGemma session as VRAM-resident.
There was a problem hiding this comment.
Update: reverted this _gpu_offload_active change in b6b6a6f. It regressed macOS/Metal (where _effective_gpu_count() is 0 but the model holds unified memory), and on a genuine CPU host training_vram never runs a real unload (it is CUDA-VRAM based and CUDA training cannot start without a GPU), so the original True is correct for every reachable case. Keeping the PR to just the UNSLOTH_ALLOW_CPU unblock.
| if self._effective_gpu_count() == 0: | ||
| env.setdefault("UNSLOTH_ALLOW_CPU", "1") |
There was a problem hiding this comment.
This condition only checks the CUDA-visible GPU count, so it is also true on supported non-CUDA accelerator hosts such as Apple Silicon/MLX and Intel XPU. In Unsloth's device detector, UNSLOTH_ALLOW_CPU=1 is checked before MLX/XPU and forces the imported process to report cuda, so the DiffusionGemma shim can lose the real backend on those hosts even though Studio explicitly supports them. Gate this on the detected Studio device being CPU, not just torch.cuda.device_count() == 0.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
get_device_type() returns mlx and xpu before it ever reads UNSLOTH_ALLOW_CPU (order: MLX, then cuda, then xpu, then the accelerator/UNSLOTH_ALLOW_CPU fallback), so on Apple Silicon/MLX and Intel XPU hosts the real backend is returned and the flag is never consulted. The diffusion shim also does not read DEVICE_TYPE at all; the visual-server binary selects its own backend (Metal/XPU/CPU) from the installed per-platform bundle. So setting UNSLOTH_ALLOW_CPU on a non-CUDA accelerator host is a no-op and cannot lose the real backend.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ba50254ec0
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| cpu_only = self._effective_gpu_count() == 0 | ||
| if cpu_only: | ||
| env.setdefault("UNSLOTH_ALLOW_CPU", "1") |
There was a problem hiding this comment.
On macOS/MLX hosts _effective_gpu_count() returns 0 because it only checks torch.cuda, but the Studio installer explicitly treats macos-arm64 as a GPU-capable llama.cpp bundle and installs the DiffusionGemma visual-server there. This branch therefore sets UNSLOTH_ALLOW_CPU and later marks _gpu_offload_active = False for a Metal-backed DiffusionGemma load, causing training_vram.free_chat_models_for_training() to skip unloading it when training starts even though it can hold unified GPU memory.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Right, that is a real regression in my previous commit. On macOS the visual server uses the Metal (macos-arm64) bundle and holds unified memory, so marking it non-resident was wrong. Reverted in b6b6a6f: _gpu_offload_active stays True, so training_vram keeps counting a Metal-backed DiffusionGemma as resident. The remaining UNSLOTH_ALLOW_CPU set is a no-op on Metal/MLX/XPU anyway (get_device_type returns mlx/xpu before reading it, and the shim does not read DEVICE_TYPE), so it only takes effect on a genuine CPU host.
…tweak; Metal hosts still hold unified memory)
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b6b6a6fcb4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # 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 self._effective_gpu_count() == 0: | ||
| env.setdefault("UNSLOTH_ALLOW_CPU", "1") |
There was a problem hiding this comment.
When Studio is installed with --no-torch/GGUF-only, install.sh installs unsloth and unsloth-zoo with --no-deps and never installs torch; in that environment this branch only sets UNSLOTH_ALLOW_CPU. The child is launched as python -m unsloth_zoo.diffusion_studio.shim, so Python imports the unsloth_zoo package before the shim, and its GPU-init path checks for torch before reaching get_device_type; UNSLOTH_ALLOW_CPU therefore does not help and the CPU DiffusionGemma runner still exits before starting. This CPU/no-torch path needs to bypass zoo GPU init rather than only allowing the CPU device.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Correct that UNSLOTH_ALLOW_CPU does not help a --no-torch/GGUF-only install, but that is out of scope here. device_type.py does if not _IS_MLX: import torch, so import unsloth_zoo (and therefore the diffusion shim) fails at import without torch, independent of this change. This PR targets the standard install (torch present, no GPU), which is verified working end to end. Making DiffusionGemma importable without torch is a separate change in unsloth_zoo.
| # 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 self._effective_gpu_count() == 0: | ||
| env.setdefault("UNSLOTH_ALLOW_CPU", "1") |
There was a problem hiding this comment.
When this added CPU-only branch succeeds on a host with torch but no GPU, _start_diffusion_server still later sets _gpu_offload_active = True before and after health. The training VRAM route treats any active GGUF with _gpu_offload_active is not False as holding VRAM (studio/backend/routes/training_vram.py:65,330), so starting training will unload the CPU-only DiffusionGemma server even though it cannot free GPU memory. Set the flag to False for this branch once healthy.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Keeping _gpu_offload_active = True here is the safe default. _effective_gpu_count() == 0 is also true on Apple/Metal and Intel XPU hosts, where the visual server holds unified/device memory; marking those non-resident would skip a needed unload and risk OOM during training (that was the regression flagged in the previous round). On a genuine CPU-only host the only effect of True is unloading before a training attempt that cannot run anyway (CUDA training needs a GPU), so there is no reachable harm.
| if self._effective_gpu_count() == 0: | ||
| env.setdefault("UNSLOTH_ALLOW_CPU", "1") |
There was a problem hiding this comment.
If a GPU machine is intentionally started with CUDA_VISIBLE_DEVICES="" to force CPU-only serving, this branch now lets the diffusion shim continue, but the command still defaults --gpu to 0; the shim's visual-engine launch sets the visual-server child's CUDA_VISIBLE_DEVICES from that argument, re-exposing GPU 0 for CUDA visual-server installs. That means a masked GPU host can unexpectedly consume GPU memory instead of staying CPU-only; keep the empty mask/CPU mode when this fallback is triggered rather than passing the default GPU ordinal.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d8565c8: when there is no visible CUDA GPU, the fallback now forces the visual-server child to stay CPU-masked (empty --gpu, so CUDA_VISIBLE_DEVICES is empty for the child), so a GPU host started with CUDA_VISIBLE_DEVICES="" no longer gets GPU 0 re-exposed via the shim default.
|
@codex review |
|
Codex Review: Didn't find any major issues. Already looking forward to the next diff. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
… GPU host does not re-expose GPU 0
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. Already looking forward to the next diff. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
1 participant
Lets DiffusionGemma GGUFs run in Studio on CPU-only hosts. Complements #6311 (which routes CPU hosts to the fork's CPU bundle that ships the DiffusionGemma binaries) by removing the last blocker on that path.
Problem
On a GPU-less host, loading a DiffusionGemma GGUF fails at load time:
The diffusion runner is launched as
python -m unsloth_zoo.diffusion_studio.shim. Importingunsloth_zoorunsdevice_type.get_device_type()at module import, which aborts when no accelerator is visible. This happens even though the shim only drives thellama-diffusion-gemma-visual-serversubprocess and does no torch GPU work itself, and even though the visual-server binary and the model run fine on CPU.Fix
In
_start_diffusion_server, when the host has no usable GPU, setUNSLOTH_ALLOW_CPU=1in the runner's child env (viasetdefault, so an explicit user value still wins). This is the same env-priming pattern already used one line above forUNSLOTH_IS_PRESENT. On GPU hosts nothing changes.Verification
Tested end-to-end CPU-only (GPUs hidden via
CUDA_VISIBLE_DEVICES="") with the fork's CPU prebuilt (b9909-mix-53618c5,linux-x64-cpu):get_device_type()error above.unsloth/diffusiongemma-26B-A4B-it-GGUF(Q4_K_M) loads on CPU (~20s), and an OpenAI-compat request with tools returns a structuredtool_callsarray (finish_reason: tool_calls,get_weather(location="San Francisco")).loaded CPU backend from .../libggml-cpu-sapphirerapids.so, everylayer N assigned to device CPU, no CUDA.Plain (non-diffusion) GGUF and GPU DiffusionGemma paths are unchanged.