fix: add Apple Silicon (MPS) support (was silently falling back to CPU) - #2
Open
gaaaork wants to merge 1 commit into
Open
fix: add Apple Silicon (MPS) support (was silently falling back to CPU)#2gaaaork wants to merge 1 commit into
gaaaork wants to merge 1 commit into
Conversation
torch.cuda.is_available() was the only device check throughout nodes.py and encoder_loader.py, so on any Mac this pack ran the 5B transformer on CPU instead of the GPU, with no warning. Confirmed on real hardware: ~1432s/step (24 min) on Apple Silicon CPU vs ~4-6s/step once routed to MPS — roughly a 250-400x difference for the exact same output. New device_compat.py centralizes the platform check (best_device(), used for picking cuda > mps > cpu; empty_cache(device), the MPS/CUDA-aware equivalent of the scattered torch.cuda.empty_cache() calls). Block-residency streaming (CUDA-specific VRAM budgeting via torch.cuda.mem_get_info()) is left CUDA-only on purpose — MPS's unified memory model doesn't need it, the existing non-CUDA fallback (whole transformer resident) already does the right thing once device correctly resolves to "mps" instead of "cpu". Tested end-to-end on Apple Silicon (M-series, 48GB unified memory, PyTorch 2.13, macOS): 4-step smoke test then a real 50-step Base generation, both via the standard Rebels SeFi Loader -> Rebels SeFi Sampler -> SaveImage graph, weight_dtype=bf16 (fp8_e4m3fn is CUDA-oriented and untested here), 1024x1024, real prompt. 50 steps completed in 349s with no errors and a correct image. No behavior change on CUDA or CPU-only machines.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Every device check in
nodes.pyandencoder_loader.pywastorch.cuda.is_available(),with no MPS branch. On a Mac, that resolves to
False, so the pack silently falls back toCPU for the whole 5B transformer — no error, no warning, just very slow generation.
Confirmed on real Apple Silicon hardware (M-series, 48GB unified memory, PyTorch 2.13):
~1432s/step (24 min) on CPU vs ~4-6s/step once routed to MPS — roughly a 250-400x
difference, same weights, same output.
Fix
Added
device_compat.py(new, small, self-contained):best_device()—cuda>mps>cpu, replaces the"cuda" if torch.cuda.is_available() else "cpu"pattern used at the one spot that actually decides where the transformer runs.empty_cache(device)— MPS/CUDA-aware cache clear, replaces the scatteredif torch.cuda.is_available(): torch.cuda.empty_cache()calls.Left untouched on purpose: the CUDA-specific VRAM-budget block-residency planner
(
_plan_residency, usingtorch.cuda.mem_get_info()). MPS's unified memory model doesn'thave a direct equivalent, and the existing "not CUDA -> keep everything resident" fallback
already does the right thing once
devicecorrectly resolves to"mps"instead of"cpu"—no new logic needed there.
Testing
Real end-to-end runs on Apple Silicon, standard Rebels SeFi Loader -> Rebels SeFi Sampler ->
SaveImage graph,
weight_dtype=bf16(fp8_e4m3fn is CUDA-oriented, not exercised here):Device: mpsin the ComfyUI log, nounsupported-op errors.
image, no errors.
No behavior change on CUDA or CPU-only machines —
best_device()returns the exact samevalue
torch.cuda.is_available()would have on those platforms.Scope
Deliberately minimal — 3 files, one new. Doesn't touch the folder-structure issue (#1),
unrelated to this fix.