From 609f3e5fa57181f440738c88bbd9172869394241 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 23 Oct 2026 14:37:50 -0600 Subject: [PATCH] Fix issue with memory manager on MPS devices --- .../diffusion_models/krea2/src/mmdit.py | 6 +++++- toolkit/memory_management/manager.py | 21 +++++++++++-------- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/extensions_built_in/diffusion_models/krea2/src/mmdit.py b/extensions_built_in/diffusion_models/krea2/src/mmdit.py index 2d66b48..5936e77 100644 --- a/extensions_built_in/diffusion_models/krea2/src/mmdit.py +++ b/extensions_built_in/diffusion_models/krea2/src/mmdit.py @@ -31,7 +31,11 @@ from torch.utils.checkpoint import checkpoint def rope(pos: Tensor, dim: int, theta: float = 1e4, ntk: float = 1.0) -> Tensor: - scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim + # MPS has no float64 + dt = torch.float32 if pos.device.type == "mps" else torch.float64 + if dt is torch.float32: + pos = pos.to(dt) + scale = torch.arange(0, dim, 2, dtype=dt, device=pos.device) / dim omega = 1.0 / ((theta * ntk) ** scale) out = torch.einsum("...n,d->...nd", pos, omega) out = torch.stack( diff --git a/toolkit/memory_management/manager.py b/toolkit/memory_management/manager.py index 5211499..c4ab070 100644 --- a/toolkit/memory_management/manager.py +++ b/toolkit/memory_management/manager.py @@ -457,15 +457,18 @@ class MemoryManager: gc.collect() torch.cuda.empty_cache() - # torch's pinned-host cache; private API, name varies by torch version - for fn_name in ("_accelerator_emptyHostCache", "_host_emptyCache"): - fn = getattr(torch._C, fn_name, None) - if fn is not None: - try: - fn() - break - except Exception: - pass + # torch's pinned-host cache; private API, name varies by torch version. + # Skipped on MPS: it segfaults there (pytorch/pytorch#197561), not catchable + on_mps = getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available() + if not on_mps: + for fn_name in ("_accelerator_emptyHostCache", "_host_emptyCache"): + fn = getattr(torch._C, fn_name, None) + if fn is not None: + try: + fn() + break + except Exception: + pass # glibc keeps freed arenas mapped; CDLL(None) resolves malloc_trim in # the running process where glibc is present and fails cleanly on