check_free_vram() computes free memory instead of used memory on the MPS branch

[ckazu]

Summary

The MPS branch of AutoTorchModule.check_free_vram() builds a tuple whose [1] - [0] is free memory, but the shared line below it treats that value as used memory. The CUDA/NPU branch is correct; only the MPS branch is inverted, so vram_limit behaves backwards on Apple Silicon.

DiffSynth-Studio/diffsynth/core/vram/layers.py

Lines 65 to 72 in b1c02ce

	def check_free_vram(self):
	if self.computation_device_type == "mps":
	gpu_mem_state = (torch.mps.current_allocated_memory(), torch.mps.recommended_max_memory())
	else:
	device = self.computation_device if not IS_NPU_AVAILABLE else get_device_name()
	gpu_mem_state = getattr(torch, self.computation_device_type).mem_get_info(device)
	used_memory = (gpu_mem_state[1] - gpu_mem_state[0]) / (1024**3)
	return used_memory < self.vram_limit

def check_free_vram(self):
    if self.computation_device_type == "mps":
        gpu_mem_state = (torch.mps.current_allocated_memory(), torch.mps.recommended_max_memory())
    else:
        device = self.computation_device if not IS_NPU_AVAILABLE else get_device_name()
        gpu_mem_state = getattr(torch, self.computation_device_type).mem_get_info(device)
    used_memory = (gpu_mem_state[1] - gpu_mem_state[0]) / (1024**3)
    return used_memory < self.vram_limit

• CUDA: torch.cuda.mem_get_info() returns (free, total), so [1] - [0] == total - free == used. Correct.
• MPS: the tuple is (current_allocated, recommended_max), so [1] - [0] == recommended_max - allocated == free. Inverted.

Reproduction

import torch
GB = 1024**3
x = torch.randn(4096, 4096, device="mps")          # ~64 MB
alloc = torch.mps.current_allocated_memory()
rmax = torch.mps.recommended_max_memory()
print(f"actually allocated       : {alloc/GB:8.3f} GB")
print(f"computed as used_memory  : {(rmax-alloc)/GB:8.3f} GB")

actually allocated       :    0.062 GB
computed as used_memory  :   51.778 GB

With 62 MB in use the function reports 51.8 GB used.

Consequence

forward() promotes a layer to the preparing device when check_free_vram() is true:

DiffSynth-Studio/diffsynth/core/vram/layers.py

Lines 546 to 547 in b1c02ce

	def forward(self, x, *args, **kwargs):
	if self.state == 1 and (self.vram_limit is None or self.check_free_vram()):

Substituting the inverted value, the predicate used_memory < vram_limit becomes

recommended_max - allocated < vram_limit   <=>   allocated > recommended_max - vram_limit

so promotion is disabled while VRAM is free and enabled once it fills up — the opposite of the intent. Setting vram_limit the way the examples do for CUDA (total - 2) reduces to allocated > 2 GB, i.e. every layer is promoted and pinned from early in the first denoising step onward, and VRAM management stops bounding anything. On a 64 GB M1 Ultra running DiffSynth-Studio/MiniMax-H3-NF4 this drove the process into swap; one run was killed by jetsam mid-generation.

The current workaround is to pass vram_limit=None, which short-circuits the predicate — but that also disables the promotion heuristic entirely rather than fixing it.

Suggested fix

Make the MPS tuple mirror mem_get_info()'s (free, total) shape so the shared arithmetic stays valid:

if self.computation_device_type == "mps":
    allocated = torch.mps.current_allocated_memory()
    recommended_max = torch.mps.recommended_max_memory()
    gpu_mem_state = (recommended_max - allocated, recommended_max)

Either that, or compute used_memory per branch instead of sharing one expression.

Environment

• macOS 26.5.1 (25F80), Mac Studio M1 Ultra, 64 GB unified memory
• Python 3.12.13, torch 2.13.0, bitsandbytes 0.50.0, transformers 5.14.1
• DiffSynth-Studio at b1c02ce (current main); the MPS branch is unchanged since it was added in support run minimax-h3 on MPS #1559
• Model: DiffSynth-Studio/MiniMax-H3-NF4 (FL2VA), device="mps"

Related

• vram_limit is not taking effect #1247 — vram_limit is not taking effect on CUDA. Different root cause, same function; worth considering together.
• Disk offload mode leaks host RAM during denoising (RSS grows every step, >15x slowdown from swap thrashing) #1564 — disk-offload host RAM growth on the same model. Unrelated code path.
• fix: Add MPS and CPU device support #1152 — open PR that guards mem_get_info() for devices lacking it, but does not touch the arithmetic; it now overlaps with the MPS branch added by support run minimax-h3 on MPS #1559.

---

_{Investigated and drafted with Claude Code. The reproduction output and the code references above were produced on the machine described in Environment.}